encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat: add Spanner vector_store_similarity_search tool

Browse Source 
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
This commit is contained in:
Google Team Member
2025-12-02 11:18:57 -08:00
committed by Copybara-Service
parent 8da61be45a
commit 090711934f
9 changed files with 948 additions and 350 deletions
Download Patch File Download Diff File Expand all files Collapse all files
contributing/samples/spanner_rag_agent/README.md
+194 -31
View File
@@ -57,9 +57,9 @@ model endpoint.
CREATE MODEL EmbeddingsModel INPUT(
content STRING(MAX),
) OUTPUT(
embeddings STRUCT<statistics STRUCT<truncated BOOL, token_count FLOAT32>, values ARRAY<FLOAT32>>,
embeddings STRUCT<values ARRAY<FLOAT32>>,
) REMOTE OPTIONS (
endpoint = '//aiplatform.googleapis.com/projects/<PROJECT_ID>/locations/us-central1/publishers/google/models/text-embedding-004'
endpoint = '//aiplatform.googleapis.com/projects/<PROJECT_ID>/locations/<LOCATION>/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="<PROJECT_ID>",
instance_id="<INSTANCE_ID>",
database_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="<PROJECT_ID>",
instance_id="<INSTANCE_ID>",
database_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,
),
],
)
```
contributing/samples/spanner_rag_agent/agent.py
+42 -196
View File
@@ -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="<PROJECT_ID>",
instance_id="<INSTANCE_ID>",
database_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 = "<PROJECT_ID>"
instance_id: str = "<INSTANCE_ID>"
database_id: str = "<DATABASE_ID>"
# 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],
)
src/google/adk/tools/spanner/search_tool.py
+253 -70
View File
@@ -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),
}
src/google/adk/tools/spanner/settings.py
+100 -2
View File
@@ -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."""
src/google/adk/tools/spanner/spanner_toolset.py
+12
View File
@@ -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
src/google/adk/tools/spanner/utils.py
+24
View File
@@ -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
tests/unittests/tools/spanner/test_search_tool.py
+230 -51
View File
@@ -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"]
tests/unittests/tools/spanner/test_spanner_tool_settings.py
+45
View File
@@ -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
)
tests/unittests/tools/spanner/test_spanner_toolset.py
+48
View File
@@ -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