From 090711934f9cd4a3327a6df626167130a217e04a Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Tue, 2 Dec 2025 11:18:23 -0800 Subject: [PATCH] feat: add Spanner vector_store_similarity_search tool The vector_store_similarity_search tool performs similarity search against data in a Spanner vector store table, using the provided Spanner tool settings for configuration. PiperOrigin-RevId: 839352057 --- .../samples/spanner_rag_agent/README.md | 225 ++++++++++-- .../samples/spanner_rag_agent/agent.py | 238 +++---------- src/google/adk/tools/spanner/search_tool.py | 323 ++++++++++++++---- src/google/adk/tools/spanner/settings.py | 102 +++++- .../adk/tools/spanner/spanner_toolset.py | 12 + src/google/adk/tools/spanner/utils.py | 24 ++ .../tools/spanner/test_search_tool.py | 281 ++++++++++++--- .../spanner/test_spanner_tool_settings.py | 45 +++ .../tools/spanner/test_spanner_toolset.py | 48 +++ 9 files changed, 948 insertions(+), 350 deletions(-) diff --git a/contributing/samples/spanner_rag_agent/README.md b/contributing/samples/spanner_rag_agent/README.md index 81489837..99b60794 100644 --- a/contributing/samples/spanner_rag_agent/README.md +++ b/contributing/samples/spanner_rag_agent/README.md @@ -57,9 +57,9 @@ model endpoint. CREATE MODEL EmbeddingsModel INPUT( content STRING(MAX), ) OUTPUT( -embeddings STRUCT, values ARRAY>, +embeddings STRUCT>, ) REMOTE OPTIONS ( -endpoint = '//aiplatform.googleapis.com/projects//locations/us-central1/publishers/google/models/text-embedding-004' +endpoint = '//aiplatform.googleapis.com/projects//locations//publishers/google/models/text-embedding-005' ); ``` @@ -187,40 +187,203 @@ type. ## Which tool to use and When? -There are a few options to perform similarity search (see the `agent.py` for -implementation details): +There are a few options to perform similarity search: -1. Wraps the built-in `similarity_search` in the Spanner Toolset. +1. Use the built-in `vector_store_similarity_search` in the Spanner Toolset with explicit `SpannerVectorStoreSettings` configuration. - - This provides an easy and controlled way to perform similarity search. - You can specify different configurations related to vector search based - on your need without having to figure out all the details for a vector - search query. + - This provides an easy way to perform similarity search. You can specify + different configurations related to vector search based on your Spanner + database vector store table setup. -2. Wraps the built-in `execute_sql` in the Spanner Toolset. + Example pseudocode (see the `agent.py` for details): - - `execute_sql` is a lower-level tool that you can have more control over - with. With the flexibility, you can specify a complicated (parameterized) - SQL query for your need, and let the `LlmAgent` pass the parameters. + ```py + from google.adk.agents.llm_agent import LlmAgent + from google.adk.tools.spanner.settings import Capabilities + from google.adk.tools.spanner.settings import SpannerToolSettings + from google.adk.tools.spanner.settings import SpannerVectorStoreSettings + from google.adk.tools.spanner.spanner_toolset import SpannerToolset -3. Use the Spanner Toolset (and all the tools that come with it) directly. + # credentials_config = SpannerCredentialsConfig(...) - - The most flexible and generic way. Instead of fixing configurations via - code, you can also specify the configurations via `instruction` to - the `LlmAgent` and let LLM to decide which tool to use and what parameters - to pass to different tools. It might even combine different tools together! - Note that in this usage, SQL generation is powered by the LlmAgent, which - can be more suitable for data analysis and assistant scenarios. - - To restrict the ability of an `LlmAgent`, `SpannerToolSet` also supports - `tool_filter` to explicitly specify allowed tools. As an example, the - following code specifies that only `execute_sql` and `get_table_schema` - are allowed: + # Define Spanner tool config with the vector store settings. + vector_store_settings = SpannerVectorStoreSettings( + project_id="", + instance_id="", + database_id="", + table_name="products", + content_column="productDescription", + embedding_column="productDescriptionEmbedding", + vector_length=768, + vertex_ai_embedding_model_name="text-embedding-005", + selected_columns=[ + "productId", + "productName", + "productDescription", + ], + nearest_neighbors_algorithm="EXACT_NEAREST_NEIGHBORS", + top_k=3, + distance_type="COSINE", + additional_filter="inventoryCount > 0", + ) - ```py - toolset = SpannerToolset( - credentials_config=credentials_config, - tool_filter=["execute_sql", "get_table_schema"], - spanner_tool_settings=SpannerToolSettings(), - ) - ``` + tool_settings = SpannerToolSettings( + capabilities=[Capabilities.DATA_READ], + vector_store_settings=vector_store_settings, + ) + # Get the Spanner toolset with the Spanner tool settings and credentials config. + spanner_toolset = SpannerToolset( + credentials_config=credentials_config, + spanner_tool_settings=tool_settings, + # Use `vector_store_similarity_search` only + tool_filter=["vector_store_similarity_search"], + ) + + root_agent = LlmAgent( + model="gemini-2.5-flash", + name="spanner_knowledge_base_agent", + description=( + "Agent to answer questions about product-specific recommendations." + ), + instruction=""" + You are a helpful assistant that answers user questions about product-specific recommendations. + 1. Always use the `vector_store_similarity_search` tool to find relevant information. + 2. If no relevant information is found, say you don't know. + 3. Present all the relevant information naturally and well formatted in your response. + """, + tools=[spanner_toolset], + ) + ``` + +2. Use the built-in `similarity_search` in the Spanner Toolset. + + - `similarity_search` is a lower-level tool, which provide the most flexible + and generic way. Specify all the necessary tool's parameters is required + when interacting with `LlmAgent` before performing the tool call. This is + more suitable for data analysis, ad-hoc query and assistant scenarios. + + Example pseudocode: + + ```py + from google.adk.agents.llm_agent import LlmAgent + from google.adk.tools.spanner.settings import Capabilities + from google.adk.tools.spanner.settings import SpannerToolSettings + from google.adk.tools.spanner.spanner_toolset import SpannerToolset + + # credentials_config = SpannerCredentialsConfig(...) + + tool_settings = SpannerToolSettings( + capabilities=[Capabilities.DATA_READ], + ) + + spanner_toolset = SpannerToolset( + credentials_config=credentials_config, + spanner_tool_settings=tool_settings, + # Use `similarity_search` only + tool_filter=["similarity_search"], + ) + + root_agent = LlmAgent( + model="gemini-2.5-flash", + name="spanner_knowledge_base_agent", + description=( + "Agent to answer questions by retrieving relevant information " + "from the Spanner database." + ), + instruction=""" + You are a helpful assistant that answers user questions to find the most relavant information from a Spanner database. + 1. Always use the `similarity_search` tool to find relevant information. + 2. If no relevant information is found, say you don't know. + 3. Present all the relevant information naturally and well formatted in your response. + """, + tools=[spanner_toolset], + ) + ``` + +3. Wraps the built-in `similarity_search` in the Spanner Toolset. + + - This provides a more controlled way to perform similarity search via code. + You can extend the tool as a wrapped function tool to have customized logic. + + Example pseudocode: + + ```py + from google.adk.agents.llm_agent import LlmAgent + + from google.adk.tools.google_tool import GoogleTool + from google.adk.tools.spanner import search_tool + import google.auth + from google.auth.credentials import Credentials + + # credentials_config = SpannerCredentialsConfig(...) + + # Create a wrapped function tool for the agent on top of the built-in + # similarity_search tool in the Spanner toolset. + # This customized tool is used to perform a Spanner KNN vector search on a + # embedded knowledge base stored in a Spanner database table. + def wrapped_spanner_similarity_search( + search_query: str, + credentials: Credentials, + ) -> str: + """Perform a similarity search on the product catalog. + + Args: + search_query: The search query to find relevant content. + + Returns: + Relevant product catalog content with sources + """ + + # ... Customized logic ... + + # Instead of fixing all parameters, you can also expose some of them for + # the LLM to decide. + return search_tool.similarity_search( + project_id="", + instance_id="", + database_id="", + table_name="products", + query=search_query, + embedding_column_to_search="productDescriptionEmbedding", + columns= [ + "productId", + "productName", + "productDescription", + ] + embedding_options={ + "vertex_ai_embedding_model_name": "text-embedding-005", + }, + credentials=credentials, + additional_filter="inventoryCount > 0", + search_options={ + "top_k": 3, + "distance_type": "EUCLIDEAN", + }, + ) + + # ... + + root_agent = LlmAgent( + model="gemini-2.5-flash", + name="spanner_knowledge_base_agent", + description=( + "Agent to answer questions about product-specific recommendations." + ), + instruction=""" + You are a helpful assistant that answers user questions about product-specific recommendations. + 1. Always use the `wrapped_spanner_similarity_search` tool to find relevant information. + 2. If no relevant information is found, say you don't know. + 3. Present all the relevant information naturally and well formatted in your response. + """, + tools=[ + # Add customized Spanner tool based on the built-in similarity_search + # in the Spanner toolset. + GoogleTool( + func=wrapped_spanner_similarity_search, + credentials_config=credentials_config, + tool_settings=tool_settings, + ), + ], + ) + ``` diff --git a/contributing/samples/spanner_rag_agent/agent.py b/contributing/samples/spanner_rag_agent/agent.py index 58ec9529..14602421 100644 --- a/contributing/samples/spanner_rag_agent/agent.py +++ b/contributing/samples/spanner_rag_agent/agent.py @@ -13,23 +13,15 @@ # limitations under the License. import os -from typing import Any -from typing import Dict -from typing import Optional from google.adk.agents.llm_agent import LlmAgent from google.adk.auth.auth_credential import AuthCredentialTypes -from google.adk.tools.base_tool import BaseTool -from google.adk.tools.google_tool import GoogleTool -from google.adk.tools.spanner import query_tool -from google.adk.tools.spanner import search_tool from google.adk.tools.spanner.settings import Capabilities from google.adk.tools.spanner.settings import SpannerToolSettings +from google.adk.tools.spanner.settings import SpannerVectorStoreSettings from google.adk.tools.spanner.spanner_credentials import SpannerCredentialsConfig -from google.adk.tools.tool_context import ToolContext +from google.adk.tools.spanner.spanner_toolset import SpannerToolset import google.auth -from google.auth.credentials import Credentials -from pydantic import BaseModel # Define an appropriate credential type # Set to None to use the application default credentials (ADC) for a quick @@ -37,9 +29,6 @@ from pydantic import BaseModel CREDENTIALS_TYPE = None -# Define Spanner tool config with read capability set to allowed. -tool_settings = SpannerToolSettings(capabilities=[Capabilities.DATA_READ]) - if CREDENTIALS_TYPE == AuthCredentialTypes.OAUTH2: # Initialize the tools to do interactive OAuth # The environment variables OAUTH_CLIENT_ID and OAUTH_CLIENT_SECRET @@ -67,172 +56,46 @@ else: credentials=application_default_credentials ) +# Follow the instructions in README.md to set up the example Spanner database. +# Replace the following settings with your specific Spanner database. -### Section 1: Extending the built-in Spanner Toolset for Custom Use Cases ### -# This example illustrates how to extend the built-in Spanner toolset to create -# a customized Spanner tool. This method is advantageous when you need to deal -# with a specific use case: -# -# 1. Streamline the end user experience by pre-configuring the tool with fixed -# parameters (such as a specific database, instance, or project) and a -# dedicated SQL query, making it perfect for a single, focused use case -# like vector search on a specific table. -# 2. Enhance functionality by adding custom logic to manage tool inputs, -# execution, and result processing, providing greater control over the -# tool's behavior. -class SpannerRagSetting(BaseModel): - """Customized Spanner RAG settings for an example use case.""" +# Define Spanner vector store settings. +vector_store_settings = SpannerVectorStoreSettings( + project_id="", + instance_id="", + database_id="", + table_name="products", + content_column="productDescription", + embedding_column="productDescriptionEmbedding", + vector_length=768, + vertex_ai_embedding_model_name="text-embedding-005", + selected_columns=[ + "productId", + "productName", + "productDescription", + ], + nearest_neighbors_algorithm="EXACT_NEAREST_NEIGHBORS", + top_k=3, + distance_type="COSINE", + additional_filter="inventoryCount > 0", +) - # Replace the following settings for your Spanner database used in the sample. - project_id: str = "" - instance_id: str = "" - database_id: str = "" +# Define Spanner tool config with the vector store settings. +tool_settings = SpannerToolSettings( + capabilities=[Capabilities.DATA_READ], + vector_store_settings=vector_store_settings, +) - # Follow the instructions in README.md, the table name is "products" and the - # Spanner embedding model name is "EmbeddingsModel" in this sample. - table_name: str = "products" - # Learn more about Spanner Vertex AI integration for embedding and Spanner - # vector search. - # https://cloud.google.com/spanner/docs/ml-tutorial-embeddings - # https://cloud.google.com/spanner/docs/vector-search/overview - embedding_model_name: str = "EmbeddingsModel" - - selected_columns: list[str] = [ - "productId", - "productName", - "productDescription", - ] - embedding_column_name: str = "productDescriptionEmbedding" - - additional_filter_expression: str = "inventoryCount > 0" - vector_distance_function: str = "EUCLIDEAN_DISTANCE" - top_k: int = 3 +# Get the Spanner toolset with the Spanner tool settings and credentials config. +# Filter the tools to only include the `vector_store_similarity_search` tool. +spanner_toolset = SpannerToolset( + credentials_config=credentials_config, + spanner_tool_settings=tool_settings, + # Comment to include all allowed tools. + tool_filter=["vector_store_similarity_search"], +) -RAG_SETTINGS = SpannerRagSetting() - - -### (Option 1) Use the built-in similarity_search tool ### -# Create a wrapped function tool for the agent on top of the built-in -# similarity_search tool in the Spanner toolset. -# This customized tool is used to perform a Spanner KNN vector search on a -# embedded knowledge base stored in a Spanner database table. -def wrapped_spanner_similarity_search( - search_query: str, - credentials: Credentials, # GoogleTool handles `credentials` automatically - settings: SpannerToolSettings, # GoogleTool handles `settings` automatically - tool_context: ToolContext, # GoogleTool handles `tool_context` automatically -) -> str: - """Perform a similarity search on the product catalog. - - Args: - search_query: The search query to find relevant content. - - Returns: - Relevant product catalog content with sources - """ - columns = RAG_SETTINGS.selected_columns.copy() - - # Instead of fixing all parameters, you can also expose some of them for - # the LLM to decide. - return search_tool.similarity_search( - RAG_SETTINGS.project_id, - RAG_SETTINGS.instance_id, - RAG_SETTINGS.database_id, - RAG_SETTINGS.table_name, - search_query, - RAG_SETTINGS.embedding_column_name, - columns, - { - "spanner_embedding_model_name": RAG_SETTINGS.embedding_model_name, - }, - credentials, - settings, - tool_context, - RAG_SETTINGS.additional_filter_expression, - { - "top_k": RAG_SETTINGS.top_k, - "distance_type": RAG_SETTINGS.vector_distance_function, - }, - ) - - -### (Option 2) Use the built-in execute_sql tool ### -# Create a wrapped function tool for the agent on top of the built-in -# execute_sql tool in the Spanner toolset. -# This customized tool is used to perform a Spanner KNN vector search on a -# embedded knowledge base stored in a Spanner database table. -# -# Compared with similarity_search, using execute_sql (a lower level tool) means -# that you have more control, but you also need to do more work (e.g. to write -# the SQL query from scratch). Consider using this option if your scenario is -# more complicated than a plain similarity search. -def wrapped_spanner_execute_sql_tool( - search_query: str, - credentials: Credentials, # GoogleTool handles `credentials` automatically - settings: SpannerToolSettings, # GoogleTool handles `settings` automatically - tool_context: ToolContext, # GoogleTool handles `tool_context` automatically -) -> str: - """Perform a similarity search on the product catalog. - - Args: - search_query: The search query to find relevant content. - - Returns: - Relevant product catalog content with sources - """ - - embedding_query = f"""SELECT embeddings.values - FROM ML.PREDICT( - MODEL {RAG_SETTINGS.embedding_model_name}, - (SELECT "{search_query}" as content) - ) - """ - - distance_alias = "distance" - columns = [f"{column}" for column in RAG_SETTINGS.selected_columns] - columns += [f"""{RAG_SETTINGS.vector_distance_function}( - {RAG_SETTINGS.embedding_column_name}, - ({embedding_query})) AS {distance_alias} - """] - columns = ", ".join(columns) - - knn_query = f""" - SELECT {columns} - FROM {RAG_SETTINGS.table_name} - WHERE {RAG_SETTINGS.additional_filter_expression} - ORDER BY {distance_alias} - LIMIT {RAG_SETTINGS.top_k} - """ - - # Customized tool based on the built-in Spanner toolset. - return query_tool.execute_sql( - project_id=RAG_SETTINGS.project_id, - instance_id=RAG_SETTINGS.instance_id, - database_id=RAG_SETTINGS.database_id, - query=knn_query, - credentials=credentials, - settings=settings, - tool_context=tool_context, - ) - - -def inspect_tool_params( - tool: BaseTool, - args: Dict[str, Any], - tool_context: ToolContext, -) -> Optional[Dict]: - """A callback function to inspect tool parameters before execution.""" - print("Inspect for tool: " + tool.name) - - actual_search_query_in_args = args.get("search_query") - # Inspect the `search_query` when calling the tool for tutorial purposes. - print(f"Tool args `search_query`: {actual_search_query_in_args}") - - pass - - -### Section 2: Create the root agent ### root_agent = LlmAgent( model="gemini-2.5-flash", name="spanner_knowledge_base_agent", @@ -241,27 +104,10 @@ root_agent = LlmAgent( ), instruction=""" You are a helpful assistant that answers user questions about product-specific recommendations. - 1. Always use the `wrapped_spanner_similarity_search` tool to find relevant information. - 2. If no relevant information is found, say you don't know. - 3. Present all the relevant information naturally and well formatted in your response. + 1. Always use the `vector_store_similarity_search` tool to find information. + 2. Directly present all the information results from the `vector_store_similarity_search` tool naturally and well formatted in your response. + 3. If no information result is returned by the `vector_store_similarity_search` tool, say you don't know. """, - tools=[ - # # (Option 1) - # # Add customized Spanner tool based on the built-in similarity_search - # # in the Spanner toolset. - GoogleTool( - func=wrapped_spanner_similarity_search, - credentials_config=credentials_config, - tool_settings=tool_settings, - ), - # # (Option 2) - # # Add customized Spanner tool based on the built-in execute_sql in - # # the Spanner toolset. - # GoogleTool( - # func=wrapped_spanner_execute_sql_tool, - # credentials_config=credentials_config, - # tool_settings=tool_settings, - # ), - ], - before_tool_callback=inspect_tool_params, + # Use the Spanner toolset for vector similarity search. + tools=[spanner_toolset], ) diff --git a/src/google/adk/tools/spanner/search_tool.py b/src/google/adk/tools/spanner/search_tool.py index b3cf797e..2b6b9777 100644 --- a/src/google/adk/tools/spanner/search_tool.py +++ b/src/google/adk/tools/spanner/search_tool.py @@ -20,23 +20,34 @@ from typing import Dict from typing import List from typing import Optional -from google.adk.tools.spanner import client -from google.adk.tools.spanner.settings import SpannerToolSettings -from google.adk.tools.tool_context import ToolContext from google.auth.credentials import Credentials from google.cloud.spanner_admin_database_v1.types import DatabaseDialect from google.cloud.spanner_v1.database import Database -# Embedding options -_SPANNER_EMBEDDING_MODEL_NAME = "spanner_embedding_model_name" -_VERTEX_AI_EMBEDDING_MODEL_ENDPOINT = "vertex_ai_embedding_model_endpoint" +from . import client +from . import utils +from .settings import APPROXIMATE_NEAREST_NEIGHBORS +from .settings import EXACT_NEAREST_NEIGHBORS +from .settings import SpannerToolSettings + +# Embedding model settings. +# Only for Spanner GoogleSQL dialect database, and use Spanner ML.PREDICT +# function. +_SPANNER_GSQL_EMBEDDING_MODEL_NAME = "spanner_googlesql_embedding_model_name" +# Only for Spanner PostgreSQL dialect database, and use spanner.ML_PREDICT_ROW +# to inferencing with Vertex AI embedding model endpoint. +_SPANNER_PG_VERTEX_AI_EMBEDDING_MODEL_ENDPOINT = ( + "spanner_postgresql_vertex_ai_embedding_model_endpoint" +) +# For both Spanner GoogleSQL and PostgreSQL dialects, use Vertex AI embedding +# model to generate embeddings for vector similarity search. +_VERTEX_AI_EMBEDDING_MODEL_NAME = "vertex_ai_embedding_model_name" +_OUTPUT_DIMENSIONALITY = "output_dimensionality" # Search options _TOP_K = "top_k" _DISTANCE_TYPE = "distance_type" _NEAREST_NEIGHBORS_ALGORITHM = "nearest_neighbors_algorithm" -_EXACT_NEAREST_NEIGHBORS = "EXACT_NEAREST_NEIGHBORS" -_APPROXIMATE_NEAREST_NEIGHBORS = "APPROXIMATE_NEAREST_NEIGHBORS" _NUM_LEAVES_TO_SEARCH = "num_leaves_to_search" # Constants @@ -48,12 +59,12 @@ _POSTGRESQL_PARAMETER_QUERY_EMBEDDING = "1" def _generate_googlesql_for_embedding_query( - spanner_embedding_model_name: str, + spanner_gsql_embedding_model_name: str, ) -> str: return f""" SELECT embeddings.values FROM ML.PREDICT( - MODEL {spanner_embedding_model_name}, + MODEL {spanner_gsql_embedding_model_name}, (SELECT CAST(@{_GOOGLESQL_PARAMETER_TEXT_QUERY} AS STRING) as content) ) """ @@ -61,37 +72,60 @@ def _generate_googlesql_for_embedding_query( def _generate_postgresql_for_embedding_query( vertex_ai_embedding_model_endpoint: str, + output_dimensionality: Optional[int], ) -> str: + instances_json = f""" + 'instances', + JSONB_BUILD_ARRAY( + JSONB_BUILD_OBJECT( + 'content', + ${_POSTGRESQL_PARAMETER_TEXT_QUERY}::TEXT + ) + ) + """ + + params_list = [] + if output_dimensionality is not None: + params_list.append(f""" + 'parameters', + JSONB_BUILD_OBJECT( + 'outputDimensionality', + {output_dimensionality} + ) + """) + + jsonb_build_args = ",\n".join([instances_json] + params_list) + return f""" - SELECT spanner.FLOAT32_ARRAY( spanner.ML_PREDICT_ROW( - '{vertex_ai_embedding_model_endpoint}', - JSONB_BUILD_OBJECT( - 'instances', - JSONB_BUILD_ARRAY( JSONB_BUILD_OBJECT( - 'content', - ${_POSTGRESQL_PARAMETER_TEXT_QUERY}::TEXT - )) + SELECT spanner.FLOAT32_ARRAY( + spanner.ML_PREDICT_ROW( + '{vertex_ai_embedding_model_endpoint}', + JSONB_BUILD_OBJECT( + {jsonb_build_args} + ) + ) -> 'predictions' -> 0 -> 'embeddings' -> 'values' ) - ) -> 'predictions'->0->'embeddings'->'values' ) """ def _get_embedding_for_query( database: Database, dialect: DatabaseDialect, - spanner_embedding_model_name: Optional[str], - vertex_ai_embedding_model_endpoint: Optional[str], + spanner_gsql_embedding_model_name: Optional[str], + spanner_pg_vertex_ai_embedding_model_endpoint: Optional[str], query: str, + output_dimensionality: Optional[int] = None, ) -> List[float]: """Gets the embedding for the query.""" if dialect == DatabaseDialect.POSTGRESQL: embedding_query = _generate_postgresql_for_embedding_query( - vertex_ai_embedding_model_endpoint + spanner_pg_vertex_ai_embedding_model_endpoint, + output_dimensionality, ) params = {f"p{_POSTGRESQL_PARAMETER_TEXT_QUERY}": query} else: embedding_query = _generate_googlesql_for_embedding_query( - spanner_embedding_model_name + spanner_gsql_embedding_model_name ) params = {_GOOGLESQL_PARAMETER_TEXT_QUERY: query} with database.snapshot() as snapshot: @@ -101,8 +135,8 @@ def _get_embedding_for_query( def _get_postgresql_distance_function(distance_type: str) -> str: return { - "COSINE_DISTANCE": "spanner.cosine_distance", - "EUCLIDEAN_DISTANCE": "spanner.euclidean_distance", + "COSINE": "spanner.cosine_distance", + "EUCLIDEAN": "spanner.euclidean_distance", "DOT_PRODUCT": "spanner.dot_product", }[distance_type] @@ -110,13 +144,13 @@ def _get_postgresql_distance_function(distance_type: str) -> str: def _get_googlesql_distance_function(distance_type: str, ann: bool) -> str: if ann: return { - "COSINE_DISTANCE": "APPROX_COSINE_DISTANCE", - "EUCLIDEAN_DISTANCE": "APPROX_EUCLIDEAN_DISTANCE", + "COSINE": "APPROX_COSINE_DISTANCE", + "EUCLIDEAN": "APPROX_EUCLIDEAN_DISTANCE", "DOT_PRODUCT": "APPROX_DOT_PRODUCT", }[distance_type] return { - "COSINE_DISTANCE": "COSINE_DISTANCE", - "EUCLIDEAN_DISTANCE": "EUCLIDEAN_DISTANCE", + "COSINE": "COSINE_DISTANCE", + "EUCLIDEAN": "EUCLIDEAN_DISTANCE", "DOT_PRODUCT": "DOT_PRODUCT", }[distance_type] @@ -172,7 +206,7 @@ def _generate_sql_for_ann( """Generates a SQL query for ANN search.""" if dialect == DatabaseDialect.POSTGRESQL: raise NotImplementedError( - f"{_APPROXIMATE_NEAREST_NEIGHBORS} is not supported for PostgreSQL" + f"{APPROXIMATE_NEAREST_NEIGHBORS} is not supported for PostgreSQL" " dialect." ) distance_function = _get_googlesql_distance_function(distance_type, ann=True) @@ -206,8 +240,6 @@ def similarity_search( columns: List[str], embedding_options: Dict[str, str], credentials: Credentials, - settings: SpannerToolSettings, - tool_context: ToolContext, additional_filter: Optional[str] = None, search_options: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: @@ -234,21 +266,34 @@ def similarity_search( columns (List[str]): A list of column names, representing the additional columns to return in the search results. embedding_options (Dict[str, str]): A dictionary of options to use for - the embedding service. The following options are supported: - - spanner_embedding_model_name: (For GoogleSQL dialect) The + the embedding service. **Exactly one of the following three keys + MUST be present in this dictionary**: + `vertex_ai_embedding_model_name`, `spanner_googlesql_embedding_model_name`, + or `spanner_postgresql_vertex_ai_embedding_model_endpoint`. + - vertex_ai_embedding_model_name (str): (Supported both **GoogleSQL and + PostgreSQL** dialects Spanner database) The name of a + public Vertex AI embedding model (e.g., `'text-embedding-005'`). + If specified, the tool generates embeddings client-side using the + Vertex AI embedding model. + - spanner_googlesql_embedding_model_name (str): (For GoogleSQL dialect) The name of the embedding model that is registered in Spanner via a `CREATE MODEL` statement. For more details, see https://cloud.google.com/spanner/docs/ml-tutorial-embeddings#generate_and_store_text_embeddings - - vertex_ai_embedding_model_endpoint: (For PostgreSQL dialect) - The fully qualified endpoint of the Vertex AI embedding model, - in the format of + If specified, embedding generation is performed using Spanner's + `ML.PREDICT` function. + - spanner_postgresql_vertex_ai_embedding_model_endpoint (str): + (For PostgreSQL dialect) The fully qualified endpoint of the Vertex AI + embedding model, in the format of `projects/$project/locations/$location/publishers/google/models/$model_name`, where $project is the project hosting the Vertex AI endpoint, $location is the location of the endpoint, and $model_name is the name of the text embedding model. + If specified, embedding generation is performed using Spanner's + `spanner.ML_PREDICT_ROW` function. + - output_dimensionality: Optional. The output dimensionality of the + embedding. If not specified, the embedding model's default output + dimensionality will be used. credentials (Credentials): The credentials to use for the request. - settings (SpannerToolSettings): The configuration for the tool. - tool_context (ToolContext): The context for the tool. additional_filter (Optional[str]): An optional filter to apply to the search query. If provided, this will be added to the WHERE clause of the final query. @@ -257,9 +302,9 @@ def similarity_search( - top_k: The number of most similar documents to return. The default value is 4. - distance_type: The distance type to use to perform the - similarity search. Valid values include "COSINE_DISTANCE", - "EUCLIDEAN_DISTANCE", and "DOT_PRODUCT". Default value is - "COSINE_DISTANCE". + similarity search. Valid values include "COSINE", + "EUCLIDEAN", and "DOT_PRODUCT". Default value is + "COSINE". - nearest_neighbors_algorithm: The nearest neighbors search algorithm to use. Valid values include "EXACT_NEAREST_NEIGHBORS" and "APPROXIMATE_NEAREST_NEIGHBORS". Default value is @@ -287,15 +332,13 @@ def similarity_search( ... embedding_column_to_search="product_description_embedding", ... columns=["product_name", "product_description", "price_in_cents"], ... credentials=credentials, - ... settings=settings, - ... tool_context=tool_context, ... additional_filter="price_in_cents < 100000", ... embedding_options={ - ... "spanner_embedding_model_name": "my_embedding_model" + ... "vertex_ai_embedding_model_name": "text-embedding-005" ... }, ... search_options={ ... "top_k": 2, - ... "distance_type": "COSINE_DISTANCE" + ... "distance_type": "COSINE" ... } ... ) { @@ -336,33 +379,68 @@ def similarity_search( embedding_options = {} if search_options is None: search_options = {} - spanner_embedding_model_name = embedding_options.get( - _SPANNER_EMBEDDING_MODEL_NAME + + exclusive_embedding_model_keys = { + _VERTEX_AI_EMBEDDING_MODEL_NAME, + _SPANNER_GSQL_EMBEDDING_MODEL_NAME, + _SPANNER_PG_VERTEX_AI_EMBEDDING_MODEL_ENDPOINT, + } + if ( + len( + exclusive_embedding_model_keys.intersection( + embedding_options.keys() + ) + ) + != 1 + ): + raise ValueError("Exactly one embedding model option must be specified.") + + vertex_ai_embedding_model_name = embedding_options.get( + _VERTEX_AI_EMBEDDING_MODEL_NAME + ) + spanner_gsql_embedding_model_name = embedding_options.get( + _SPANNER_GSQL_EMBEDDING_MODEL_NAME + ) + spanner_pg_vertex_ai_embedding_model_endpoint = embedding_options.get( + _SPANNER_PG_VERTEX_AI_EMBEDDING_MODEL_ENDPOINT ) if ( database.database_dialect == DatabaseDialect.GOOGLE_STANDARD_SQL - and spanner_embedding_model_name is None + and vertex_ai_embedding_model_name is None + and spanner_gsql_embedding_model_name is None ): raise ValueError( - f"embedding_options['{_SPANNER_EMBEDDING_MODEL_NAME}']" - " must be specified for GoogleSQL dialect." + f"embedding_options['{_VERTEX_AI_EMBEDDING_MODEL_NAME}'] or" + f" embedding_options['{_SPANNER_GSQL_EMBEDDING_MODEL_NAME}'] must be" + " specified for GoogleSQL dialect Spanner database." ) - vertex_ai_embedding_model_endpoint = embedding_options.get( - _VERTEX_AI_EMBEDDING_MODEL_ENDPOINT - ) if ( database.database_dialect == DatabaseDialect.POSTGRESQL - and vertex_ai_embedding_model_endpoint is None + and vertex_ai_embedding_model_name is None + and spanner_pg_vertex_ai_embedding_model_endpoint is None ): raise ValueError( - f"embedding_options['{_VERTEX_AI_EMBEDDING_MODEL_ENDPOINT}']" - " must be specified for PostgreSQL dialect." + f"embedding_options['{_VERTEX_AI_EMBEDDING_MODEL_NAME}'] or" + f" embedding_options['{_SPANNER_PG_VERTEX_AI_EMBEDDING_MODEL_ENDPOINT}']" + " must be specified for PostgreSQL dialect Spanner database." + ) + output_dimensionality = embedding_options.get(_OUTPUT_DIMENSIONALITY) + if ( + output_dimensionality is not None + and spanner_gsql_embedding_model_name is not None + ): + # Currently, Spanner GSQL Model ML.PREDICT does not support + # output_dimensionality parameter for inference embedding models. + raise ValueError( + f"embedding_options[{_OUTPUT_DIMENSIONALITY}] is not supported when" + f" embedding_options['{_SPANNER_GSQL_EMBEDDING_MODEL_NAME}'] is" + " specified." ) # Use cosine distance by default. distance_type = search_options.get(_DISTANCE_TYPE) if distance_type is None: - distance_type = "COSINE_DISTANCE" + distance_type = "COSINE" top_k = search_options.get(_TOP_K) if top_k is None: @@ -370,26 +448,36 @@ def similarity_search( # Use EXACT_NEAREST_NEIGHBORS (i.e. kNN) by default. nearest_neighbors_algorithm = search_options.get( - _NEAREST_NEIGHBORS_ALGORITHM, _EXACT_NEAREST_NEIGHBORS + _NEAREST_NEIGHBORS_ALGORITHM, + EXACT_NEAREST_NEIGHBORS, ) if nearest_neighbors_algorithm not in ( - _EXACT_NEAREST_NEIGHBORS, - _APPROXIMATE_NEAREST_NEIGHBORS, + EXACT_NEAREST_NEIGHBORS, + APPROXIMATE_NEAREST_NEIGHBORS, ): raise NotImplementedError( f"Unsupported search_options['{_NEAREST_NEIGHBORS_ALGORITHM}']:" f" {nearest_neighbors_algorithm}" ) - embedding = _get_embedding_for_query( - database, - database.database_dialect, - spanner_embedding_model_name, - vertex_ai_embedding_model_endpoint, - query, - ) + # Generate embedding for the query according to the embedding options. + if vertex_ai_embedding_model_name: + embedding = utils.embed_contents( + vertex_ai_embedding_model_name, + [query], + output_dimensionality, + )[0] + else: + embedding = _get_embedding_for_query( + database, + database.database_dialect, + spanner_gsql_embedding_model_name, + spanner_pg_vertex_ai_embedding_model_endpoint, + query, + output_dimensionality, + ) - if nearest_neighbors_algorithm == _EXACT_NEAREST_NEIGHBORS: + if nearest_neighbors_algorithm == EXACT_NEAREST_NEIGHBORS: sql = _generate_sql_for_knn( database.database_dialect, table_name, @@ -438,5 +526,100 @@ def similarity_search( except Exception as ex: return { "status": "ERROR", - "error_details": str(ex), + "error_details": repr(ex), + } + + +def vector_store_similarity_search( + query: str, + credentials: Credentials, + settings: SpannerToolSettings, +) -> Dict[str, Any]: + """Performs a semantic similarity search to retrieve relevant context from the Spanner vector store. + + This function performs vector similarity search directly on a vector store + table in Spanner database and returns the relevant data. + + Args: + query (str): The search string based on the user's question. + credentials (Credentials): The credentials to use for the request. + settings (SpannerToolSettings): The configuration for the tool. + + Returns: + Dict[str, Any]: A dictionary representing the result of the search. + On success, it contains {"status": "SUCCESS", "rows": [...]}. The last + column of each row is the distance between the query and the row result. + On error, it contains {"status": "ERROR", "error_details": "..."}. + + Examples: + >>> vector_store_similarity_search( + ... query="Spanner database optimization techniques for high QPS", + ... credentials=credentials, + ... settings=settings + ... ) + { + "status": "SUCCESS", + "rows": [ + ( + "Optimizing Query Performance", + 0.12, + ), + ( + "Schema Design Best Practices", + 0.25, + ), + ( + "Using Secondary Indexes Effectively", + 0.31, + ), + ... + ], + } + """ + + try: + if not settings or not settings.vector_store_settings: + raise ValueError("Spanner vector store settings are not set.") + + # Get the embedding model settings. + embedding_options = { + _VERTEX_AI_EMBEDDING_MODEL_NAME: ( + settings.vector_store_settings.vertex_ai_embedding_model_name + ), + _OUTPUT_DIMENSIONALITY: settings.vector_store_settings.vector_length, + } + + # Get the search settings. + search_options = { + _TOP_K: settings.vector_store_settings.top_k, + _DISTANCE_TYPE: settings.vector_store_settings.distance_type, + _NEAREST_NEIGHBORS_ALGORITHM: ( + settings.vector_store_settings.nearest_neighbors_algorithm + ), + } + if ( + settings.vector_store_settings.nearest_neighbors_algorithm + == APPROXIMATE_NEAREST_NEIGHBORS + ): + search_options[_NUM_LEAVES_TO_SEARCH] = ( + settings.vector_store_settings.num_leaves_to_search + ) + + return similarity_search( + project_id=settings.vector_store_settings.project_id, + instance_id=settings.vector_store_settings.instance_id, + database_id=settings.vector_store_settings.database_id, + table_name=settings.vector_store_settings.table_name, + query=query, + embedding_column_to_search=settings.vector_store_settings.embedding_column, + columns=settings.vector_store_settings.selected_columns, + embedding_options=embedding_options, + credentials=credentials, + additional_filter=settings.vector_store_settings.additional_filter, + search_options=search_options, + ) + except Exception as ex: + return { + "status": "ERROR", + "error_details": repr(ex), } diff --git a/src/google/adk/tools/spanner/settings.py b/src/google/adk/tools/spanner/settings.py index 5d097258..a76331ba 100644 --- a/src/google/adk/tools/spanner/settings.py +++ b/src/google/adk/tools/spanner/settings.py @@ -16,20 +16,115 @@ from __future__ import annotations from enum import Enum from typing import List +from typing import Literal +from typing import Optional from pydantic import BaseModel +from pydantic import model_validator from ...utils.feature_decorator import experimental +# Vector similarity search nearest neighbors search algorithms. +EXACT_NEAREST_NEIGHBORS = "EXACT_NEAREST_NEIGHBORS" +APPROXIMATE_NEAREST_NEIGHBORS = "APPROXIMATE_NEAREST_NEIGHBORS" +NearestNeighborsAlgorithm = Literal[ + EXACT_NEAREST_NEIGHBORS, + APPROXIMATE_NEAREST_NEIGHBORS, +] + class Capabilities(Enum): """Capabilities indicating what type of operation tools are allowed to be performed on Spanner.""" - DATA_READ = 'data_read' + DATA_READ = "data_read" """Read only data operations tools are allowed.""" -@experimental('Tool settings defaults may have breaking change in the future.') +class SpannerVectorStoreSettings(BaseModel): + """Settings for Spanner Vector Store. + + This is used for vector similarity search in a Spanner vector store table. + Provide the vector store table and the embedding model settings to use with + the `vector_store_similarity_search` tool. + """ + + project_id: str + """Required. The GCP project id in which the Spanner database resides.""" + + instance_id: str + """Required. The instance id of the Spanner database.""" + + database_id: str + """Required. The database id of the Spanner database.""" + + table_name: str + """Required. The name of the vector store table to use for vector similarity search.""" + + content_column: str + """Required. The name of the content column in the vector store table. By default, this column value is also returned as part of the vector similarity search result.""" + + embedding_column: str + """Required. The name of the embedding column to search in the vector store table.""" + + vector_length: int + """Required. The the dimension of the vectors in the `embedding_column`.""" + + vertex_ai_embedding_model_name: str + """Required. The Vertex AI embedding model name, which is used to generate embeddings for vector store and vector similarity search. + For example, 'text-embedding-005'. + + Note: the output dimensionality of the embedding model should be the same as the value specified in the `vector_length` field. + Otherwise, a runtime error might be raised during a query. + """ + + selected_columns: List[str] = [] + """Required. The vector store table columns to return in the vector similarity search result. + + By default, only the `content_column` value and the distance value are returned. + If sepecified, the list of selected columns and the distance value are returned. + For example, if `selected_columns` is ['col1', 'col2'], then the result will contain the values of 'col1' and 'col2' columns and the distance value. + """ + + nearest_neighbors_algorithm: NearestNeighborsAlgorithm = ( + "EXACT_NEAREST_NEIGHBORS" + ) + """The algorithm used to perform vector similarity search. This value can be EXACT_NEAREST_NEIGHBORS or APPROXIMATE_NEAREST_NEIGHBORS. + + For more details about EXACT_NEAREST_NEIGHBORS, see https://docs.cloud.google.com/spanner/docs/find-k-nearest-neighbors + For more details about APPROXIMATE_NEAREST_NEIGHBORS, see https://docs.cloud.google.com/spanner/docs/find-approximate-nearest-neighbors + """ + + top_k: int = 4 + """Required. The number of neighbors to return for each vector similarity search query. The default value is 4.""" + + distance_type: str = "COSINE" + """Required. The distance metric used to build the vector index or perform vector similarity search. This value can be COSINE, DOT_PRODUCT, or EUCLIDEAN.""" + + num_leaves_to_search: Optional[int] = None + """Optional. This option specifies how many leaf nodes of the index are searched. + + Note: this option is only used when the nearest neighbors search algorithm (`nearest_neighbors_algorithm`) is APPROXIMATE_NEAREST_NEIGHBORS. + For more details, see https://docs.cloud.google.com/spanner/docs/vector-index-best-practices + """ + + additional_filter: Optional[str] = None + """Optional. An optional filter to apply to the search query. If provided, this will be added to the WHERE clause of the final query.""" + + @model_validator(mode="after") + def __post_init__(self): + """Validate the embedding settings.""" + if not self.vector_length or self.vector_length <= 0: + raise ValueError( + "Invalid vector length in the Spanner vector store settings." + ) + + if not self.selected_columns: + self.selected_columns = [self.content_column] + + return self + + +@experimental("Tool settings defaults may have breaking change in the future.") class SpannerToolSettings(BaseModel): """Settings for Spanner tools.""" @@ -44,3 +139,6 @@ class SpannerToolSettings(BaseModel): max_executed_query_result_rows: int = 50 """Maximum number of rows to return from a query result.""" + + vector_store_settings: Optional[SpannerVectorStoreSettings] = None + """Settings for Spanner vector store and vector similarity search.""" diff --git a/src/google/adk/tools/spanner/spanner_toolset.py b/src/google/adk/tools/spanner/spanner_toolset.py index 861314ab..6496014f 100644 --- a/src/google/adk/tools/spanner/spanner_toolset.py +++ b/src/google/adk/tools/spanner/spanner_toolset.py @@ -47,6 +47,8 @@ class SpannerToolset(BaseToolset): - spanner_list_named_schemas - spanner_get_table_schema - spanner_execute_sql + - spanner_similarity_search + - spanner_vector_store_similarity_search """ def __init__( @@ -121,6 +123,16 @@ class SpannerToolset(BaseToolset): tool_settings=self._tool_settings, ) ) + if self._tool_settings.vector_store_settings: + # Only add the vector store similarity search tool if the vector store + # settings are specified. + all_tools.append( + GoogleTool( + func=search_tool.vector_store_similarity_search, + credentials_config=self._credentials_config, + tool_settings=self._tool_settings, + ) + ) return [ tool diff --git a/src/google/adk/tools/spanner/utils.py b/src/google/adk/tools/spanner/utils.py index 64c03859..f1b710ec 100644 --- a/src/google/adk/tools/spanner/utils.py +++ b/src/google/adk/tools/spanner/utils.py @@ -105,3 +105,27 @@ def execute_sql( "status": "ERROR", "error_details": str(ex), } + + +def embed_contents( + vertex_ai_embedding_model_name: str, + contents: list[str], + output_dimensionality: Optional[int] = None, +) -> list[list[float]]: + """Embed the given contents into list of vectors using the Vertex AI embedding model endpoint.""" + try: + from google.genai import Client + from google.genai.types import EmbedContentConfig + + client = Client() + config = EmbedContentConfig() + if output_dimensionality: + config.output_dimensionality = output_dimensionality + response = client.models.embed_content( + model=vertex_ai_embedding_model_name, + contents=contents, + config=config, + ) + return [list(e.values) for e in response.embeddings] + except Exception as ex: + raise RuntimeError(f"Failed to embed content: {ex!r}") from ex diff --git a/tests/unittests/tools/spanner/test_search_tool.py b/tests/unittests/tools/spanner/test_search_tool.py index 1b330b01..c69aa444 100644 --- a/tests/unittests/tools/spanner/test_search_tool.py +++ b/tests/unittests/tools/spanner/test_search_tool.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest import mock from unittest.mock import MagicMock -from unittest.mock import patch +from google.adk.tools.spanner import client from google.adk.tools.spanner import search_tool +from google.adk.tools.spanner import utils from google.cloud.spanner_admin_database_v1.types import DatabaseDialect import pytest @@ -35,29 +37,59 @@ def mock_spanner_ids(): } -@patch("google.adk.tools.spanner.client.get_spanner_client") +@pytest.mark.parametrize( + ("embedding_option_key", "embedding_option_value", "expected_embedding"), + [ + pytest.param( + "spanner_googlesql_embedding_model_name", + "EmbeddingsModel", + [0.1, 0.2, 0.3], + id="spanner_googlesql_embedding_model", + ), + pytest.param( + "vertex_ai_embedding_model_name", + "text-embedding-005", + [0.4, 0.5, 0.6], + id="vertex_ai_embedding_model", + ), + ], +) +@mock.patch.object(utils, "embed_contents") +@mock.patch.object(client, "get_spanner_client") def test_similarity_search_knn_success( - mock_get_spanner_client, mock_spanner_ids, mock_credentials + mock_get_spanner_client, + mock_embed_contents, + mock_spanner_ids, + mock_credentials, + embedding_option_key, + embedding_option_value, + expected_embedding, ): """Test similarity_search function with kNN success.""" mock_spanner_client = MagicMock() mock_instance = MagicMock() mock_database = MagicMock() mock_snapshot = MagicMock() - mock_embedding_result = MagicMock() - mock_embedding_result.one.return_value = ([0.1, 0.2, 0.3],) - # First call to execute_sql is for getting the embedding - # Second call is for the kNN search - mock_snapshot.execute_sql.side_effect = [ - mock_embedding_result, - iter([("result1",), ("result2",)]), - ] mock_database.snapshot.return_value.__enter__.return_value = mock_snapshot mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL mock_instance.database.return_value = mock_database mock_spanner_client.instance.return_value = mock_instance mock_get_spanner_client.return_value = mock_spanner_client + if embedding_option_key == "vertex_ai_embedding_model_name": + mock_embed_contents.return_value = [expected_embedding] + # execute_sql is called once for the kNN search + mock_snapshot.execute_sql.return_value = iter([("result1",), ("result2",)]) + else: + mock_embedding_result = MagicMock() + mock_embedding_result.one.return_value = (expected_embedding,) + # First call to execute_sql is for getting the embedding, + # second call is for the kNN search + mock_snapshot.execute_sql.side_effect = [ + mock_embedding_result, + iter([("result1",), ("result2",)]), + ] + result = search_tool.similarity_search( project_id=mock_spanner_ids["project_id"], instance_id=mock_spanner_ids["instance_id"], @@ -66,10 +98,8 @@ def test_similarity_search_knn_success( query="test query", embedding_column_to_search="embedding_col", columns=["col1"], - embedding_options={"spanner_embedding_model_name": "test_model"}, + embedding_options={embedding_option_key: embedding_option_value}, credentials=mock_credentials, - settings=MagicMock(), - tool_context=MagicMock(), ) assert result["status"] == "SUCCESS", result assert result["rows"] == [("result1",), ("result2",)] @@ -79,10 +109,14 @@ def test_similarity_search_knn_success( sql = call_args.args[0] assert "COSINE_DISTANCE" in sql assert "@embedding" in sql - assert call_args.kwargs == {"params": {"embedding": [0.1, 0.2, 0.3]}} + assert call_args.kwargs == {"params": {"embedding": expected_embedding}} + if embedding_option_key == "vertex_ai_embedding_model_name": + mock_embed_contents.assert_called_once_with( + embedding_option_value, ["test query"], None + ) -@patch("google.adk.tools.spanner.client.get_spanner_client") +@mock.patch.object(client, "get_spanner_client") def test_similarity_search_ann_success( mock_get_spanner_client, mock_spanner_ids, mock_credentials ): @@ -113,10 +147,10 @@ def test_similarity_search_ann_success( query="test query", embedding_column_to_search="embedding_col", columns=["col1"], - embedding_options={"spanner_embedding_model_name": "test_model"}, + embedding_options={ + "spanner_googlesql_embedding_model_name": "test_model" + }, credentials=mock_credentials, - settings=MagicMock(), - tool_context=MagicMock(), search_options={ "nearest_neighbors_algorithm": "APPROXIMATE_NEAREST_NEIGHBORS" }, @@ -130,7 +164,7 @@ def test_similarity_search_ann_success( assert call_args.kwargs == {"params": {"embedding": [0.1, 0.2, 0.3]}} -@patch("google.adk.tools.spanner.client.get_spanner_client") +@mock.patch.object(client, "get_spanner_client") def test_similarity_search_error( mock_get_spanner_client, mock_spanner_ids, mock_credentials ): @@ -143,17 +177,17 @@ def test_similarity_search_error( table_name=mock_spanner_ids["table_name"], query="test query", embedding_column_to_search="embedding_col", - embedding_options={"spanner_embedding_model_name": "test_model"}, + embedding_options={ + "spanner_googlesql_embedding_model_name": "test_model" + }, columns=["col1"], credentials=mock_credentials, - settings=MagicMock(), - tool_context=MagicMock(), ) assert result["status"] == "ERROR" - assert result["error_details"] == "Test Exception" + assert "Test Exception" in result["error_details"] -@patch("google.adk.tools.spanner.client.get_spanner_client") +@mock.patch.object(client, "get_spanner_client") def test_similarity_search_postgresql_knn_success( mock_get_spanner_client, mock_spanner_ids, mock_credentials ): @@ -182,10 +216,12 @@ def test_similarity_search_postgresql_knn_success( query="test query", embedding_column_to_search="embedding_col", columns=["col1"], - embedding_options={"vertex_ai_embedding_model_endpoint": "test_endpoint"}, + embedding_options={ + "spanner_postgresql_vertex_ai_embedding_model_endpoint": ( + "test_endpoint" + ) + }, credentials=mock_credentials, - settings=MagicMock(), - tool_context=MagicMock(), ) assert result["status"] == "SUCCESS", result assert result["rows"] == [("pg_result",)] @@ -196,7 +232,7 @@ def test_similarity_search_postgresql_knn_success( assert call_args.kwargs == {"params": {"p1": [0.1, 0.2, 0.3]}} -@patch("google.adk.tools.spanner.client.get_spanner_client") +@mock.patch.object(client, "get_spanner_client") def test_similarity_search_postgresql_ann_unsupported( mock_get_spanner_client, mock_spanner_ids, mock_credentials ): @@ -217,27 +253,28 @@ def test_similarity_search_postgresql_ann_unsupported( query="test query", embedding_column_to_search="embedding_col", columns=["col1"], - embedding_options={"vertex_ai_embedding_model_endpoint": "test_endpoint"}, + embedding_options={ + "spanner_postgresql_vertex_ai_embedding_model_endpoint": ( + "test_endpoint" + ) + }, credentials=mock_credentials, - settings=MagicMock(), - tool_context=MagicMock(), search_options={ "nearest_neighbors_algorithm": "APPROXIMATE_NEAREST_NEIGHBORS" }, ) assert result["status"] == "ERROR" assert ( - result["error_details"] - == "APPROXIMATE_NEAREST_NEIGHBORS is not supported for PostgreSQL" - " dialect." + "APPROXIMATE_NEAREST_NEIGHBORS is not supported for PostgreSQL dialect." + in result["error_details"] ) -@patch("google.adk.tools.spanner.client.get_spanner_client") -def test_similarity_search_missing_spanner_embedding_model_name_error( +@mock.patch.object(client, "get_spanner_client") +def test_similarity_search_gsql_missing_embedding_model_error( mock_get_spanner_client, mock_spanner_ids, mock_credentials ): - """Test similarity_search with missing spanner_embedding_model_name.""" + """Test similarity_search with missing embedding_options for GoogleSQL dialect.""" mock_spanner_client = MagicMock() mock_instance = MagicMock() mock_database = MagicMock() @@ -254,24 +291,27 @@ def test_similarity_search_missing_spanner_embedding_model_name_error( query="test query", embedding_column_to_search="embedding_col", columns=["col1"], - embedding_options={}, + embedding_options={ + "spanner_postgresql_vertex_ai_embedding_model_endpoint": ( + "test_endpoint" + ) + }, credentials=mock_credentials, - settings=MagicMock(), - tool_context=MagicMock(), ) assert result["status"] == "ERROR" assert ( - "embedding_options['spanner_embedding_model_name'] must be" - " specified for GoogleSQL dialect." + "embedding_options['vertex_ai_embedding_model_name'] or" + " embedding_options['spanner_googlesql_embedding_model_name'] must be" + " specified for GoogleSQL dialect Spanner database." in result["error_details"] ) -@patch("google.adk.tools.spanner.client.get_spanner_client") -def test_similarity_search_missing_vertex_ai_embedding_model_endpoint_error( +@mock.patch.object(client, "get_spanner_client") +def test_similarity_search_pg_missing_embedding_model_error( mock_get_spanner_client, mock_spanner_ids, mock_credentials ): - """Test similarity_search with missing vertex_ai_embedding_model_endpoint.""" + """Test similarity_search with missing embedding_options for PostgreSQL dialect.""" mock_spanner_client = MagicMock() mock_instance = MagicMock() mock_database = MagicMock() @@ -288,14 +328,153 @@ def test_similarity_search_missing_vertex_ai_embedding_model_endpoint_error( query="test query", embedding_column_to_search="embedding_col", columns=["col1"], - embedding_options={}, + embedding_options={ + "spanner_googlesql_embedding_model_name": "EmbeddingsModel" + }, credentials=mock_credentials, - settings=MagicMock(), - tool_context=MagicMock(), ) assert result["status"] == "ERROR" assert ( - "embedding_options['vertex_ai_embedding_model_endpoint'] must " - "be specified for PostgreSQL dialect." + "embedding_options['vertex_ai_embedding_model_name'] or" + " embedding_options['spanner_postgresql_vertex_ai_embedding_model_endpoint']" + " must be specified for PostgreSQL dialect Spanner database." in result["error_details"] ) + + +@pytest.mark.parametrize( + "embedding_options", + [ + pytest.param( + { + "vertex_ai_embedding_model_name": "test-model", + "spanner_googlesql_embedding_model_name": "test-model-2", + }, + id="vertex_ai_and_googlesql", + ), + pytest.param( + { + "vertex_ai_embedding_model_name": "test-model", + "spanner_postgresql_vertex_ai_embedding_model_endpoint": ( + "test-endpoint" + ), + }, + id="vertex_ai_and_postgresql", + ), + pytest.param( + { + "spanner_googlesql_embedding_model_name": "test-model", + "spanner_postgresql_vertex_ai_embedding_model_endpoint": ( + "test-endpoint" + ), + }, + id="googlesql_and_postgresql", + ), + pytest.param( + { + "vertex_ai_embedding_model_name": "test-model", + "spanner_googlesql_embedding_model_name": "test-model-2", + "spanner_postgresql_vertex_ai_embedding_model_endpoint": ( + "test-endpoint" + ), + }, + id="all_three_models", + ), + pytest.param( + {}, + id="no_models", + ), + ], +) +@mock.patch.object(client, "get_spanner_client") +def test_similarity_search_multiple_embedding_options_error( + mock_get_spanner_client, + mock_spanner_ids, + mock_credentials, + embedding_options, +): + """Test similarity_search with multiple embedding models.""" + mock_spanner_client = MagicMock() + mock_instance = MagicMock() + mock_database = MagicMock() + mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL + mock_instance.database.return_value = mock_database + mock_spanner_client.instance.return_value = mock_instance + mock_get_spanner_client.return_value = mock_spanner_client + + result = search_tool.similarity_search( + project_id=mock_spanner_ids["project_id"], + instance_id=mock_spanner_ids["instance_id"], + database_id=mock_spanner_ids["database_id"], + table_name=mock_spanner_ids["table_name"], + query="test query", + embedding_column_to_search="embedding_col", + columns=["col1"], + embedding_options=embedding_options, + credentials=mock_credentials, + ) + assert result["status"] == "ERROR" + assert ( + "Exactly one embedding model option must be specified." + in result["error_details"] + ) + + +@mock.patch.object(client, "get_spanner_client") +def test_similarity_search_output_dimensionality_gsql_error( + mock_get_spanner_client, mock_spanner_ids, mock_credentials +): + """Test similarity_search with output_dimensionality and spanner_googlesql_embedding_model_name.""" + mock_spanner_client = MagicMock() + mock_instance = MagicMock() + mock_database = MagicMock() + mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL + mock_instance.database.return_value = mock_database + mock_spanner_client.instance.return_value = mock_instance + mock_get_spanner_client.return_value = mock_spanner_client + + result = search_tool.similarity_search( + project_id=mock_spanner_ids["project_id"], + instance_id=mock_spanner_ids["instance_id"], + database_id=mock_spanner_ids["database_id"], + table_name=mock_spanner_ids["table_name"], + query="test query", + embedding_column_to_search="embedding_col", + columns=["col1"], + embedding_options={ + "spanner_googlesql_embedding_model_name": "EmbeddingsModel", + "output_dimensionality": 128, + }, + credentials=mock_credentials, + ) + assert result["status"] == "ERROR" + assert "is not supported when" in result["error_details"] + + +@mock.patch.object(client, "get_spanner_client") +def test_similarity_search_unsupported_algorithm_error( + mock_get_spanner_client, mock_spanner_ids, mock_credentials +): + """Test similarity_search with an unsupported nearest neighbors algorithm.""" + mock_spanner_client = MagicMock() + mock_instance = MagicMock() + mock_database = MagicMock() + mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL + mock_instance.database.return_value = mock_database + mock_spanner_client.instance.return_value = mock_instance + mock_get_spanner_client.return_value = mock_spanner_client + + result = search_tool.similarity_search( + project_id=mock_spanner_ids["project_id"], + instance_id=mock_spanner_ids["instance_id"], + database_id=mock_spanner_ids["database_id"], + table_name=mock_spanner_ids["table_name"], + query="test query", + embedding_column_to_search="embedding_col", + columns=["col1"], + embedding_options={"vertex_ai_embedding_model_name": "test-model"}, + credentials=mock_credentials, + search_options={"nearest_neighbors_algorithm": "INVALID_ALGORITHM"}, + ) + assert result["status"] == "ERROR" + assert "Unsupported search_options" in result["error_details"] diff --git a/tests/unittests/tools/spanner/test_spanner_tool_settings.py b/tests/unittests/tools/spanner/test_spanner_tool_settings.py index f74922b2..730c9e0e 100644 --- a/tests/unittests/tools/spanner/test_spanner_tool_settings.py +++ b/tests/unittests/tools/spanner/test_spanner_tool_settings.py @@ -15,9 +15,23 @@ from __future__ import annotations from google.adk.tools.spanner.settings import SpannerToolSettings +from google.adk.tools.spanner.settings import SpannerVectorStoreSettings +from pydantic import ValidationError import pytest +def common_spanner_vector_store_settings(vector_length=None): + return { + "project_id": "test-project", + "instance_id": "test-instance", + "database_id": "test-database", + "table_name": "test-table", + "content_column": "test-content-column", + "embedding_column": "test-embedding-column", + "vector_length": 128 if vector_length is None else vector_length, + } + + def test_spanner_tool_settings_experimental_warning(): """Test SpannerToolSettings experimental warning.""" with pytest.warns( @@ -25,3 +39,34 @@ def test_spanner_tool_settings_experimental_warning(): match="Tool settings defaults may have breaking change in the future.", ): SpannerToolSettings() + + +def test_spanner_vector_store_settings_all_fields_present(): + """Test SpannerVectorStoreSettings with all required fields present.""" + settings = SpannerVectorStoreSettings( + **common_spanner_vector_store_settings(), + vertex_ai_embedding_model_name="test-embedding-model", + ) + assert settings is not None + assert settings.selected_columns == ["test-content-column"] + assert settings.vertex_ai_embedding_model_name == "test-embedding-model" + + +def test_spanner_vector_store_settings_missing_embedding_model_name(): + """Test SpannerVectorStoreSettings with missing vertex_ai_embedding_model_name.""" + with pytest.raises(ValidationError) as excinfo: + SpannerVectorStoreSettings(**common_spanner_vector_store_settings()) + assert "Field required" in str(excinfo.value) + assert "vertex_ai_embedding_model_name" in str(excinfo.value) + + +def test_spanner_vector_store_settings_invalid_vector_length(): + """Test SpannerVectorStoreSettings with invalid vector_length.""" + with pytest.raises(ValidationError) as excinfo: + SpannerVectorStoreSettings( + **common_spanner_vector_store_settings(vector_length=0), + vertex_ai_embedding_model_name="test-embedding-model", + ) + assert "Invalid vector length in the Spanner vector store settings." in str( + excinfo.value + ) diff --git a/tests/unittests/tools/spanner/test_spanner_toolset.py b/tests/unittests/tools/spanner/test_spanner_toolset.py index 16383255..a583a2f8 100644 --- a/tests/unittests/tools/spanner/test_spanner_toolset.py +++ b/tests/unittests/tools/spanner/test_spanner_toolset.py @@ -18,6 +18,7 @@ from google.adk.tools.google_tool import GoogleTool from google.adk.tools.spanner import SpannerCredentialsConfig from google.adk.tools.spanner import SpannerToolset from google.adk.tools.spanner.settings import SpannerToolSettings +from google.adk.tools.spanner.settings import SpannerVectorStoreSettings import pytest @@ -184,3 +185,50 @@ async def test_spanner_toolset_without_read_capability( expected_tool_names = set(returned_tools) actual_tool_names = set([tool.name for tool in tools]) assert actual_tool_names == expected_tool_names + + +@pytest.mark.asyncio +async def test_spanner_toolset_with_vector_store_search(): + """Test Spanner toolset with vector store search. + + This test verifies the behavior of the Spanner toolset when vector store + settings is provided. + """ + credentials_config = SpannerCredentialsConfig( + client_id="abc", client_secret="def" + ) + + spanner_tool_settings = SpannerToolSettings( + vector_store_settings=SpannerVectorStoreSettings( + project_id="test-project", + instance_id="test-instance", + database_id="test-database", + table_name="test-table", + content_column="test-content-column", + embedding_column="test-embedding-column", + vector_length=128, + vertex_ai_embedding_model_name="test-embedding-model", + ) + ) + toolset = SpannerToolset( + credentials_config=credentials_config, + spanner_tool_settings=spanner_tool_settings, + ) + tools = await toolset.get_tools() + assert tools is not None + + assert len(tools) == 8 + assert all([isinstance(tool, GoogleTool) for tool in tools]) + + expected_tool_names = set([ + "list_table_names", + "list_table_indexes", + "list_table_index_columns", + "list_named_schemas", + "get_table_schema", + "execute_sql", + "similarity_search", + "vector_store_similarity_search", + ]) + actual_tool_names = set([tool.name for tool in tools]) + assert actual_tool_names == expected_tool_names