encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-03-30 10:57:20 -07:00
Code Issues Packages Projects Releases Wiki Activity

chore: Sample Spanner RAG agent that wraps search_tool

Browse Source 
Also modified README to add instructions on when to use which tool.

PiperOrigin-RevId: 810563458
This commit is contained in:
Google Team Member
2025-09-23 13:37:40 -07:00
committed by Copybara-Service
parent 657369cffe
commit cbb609233b
2 changed files with 118 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
contributing/samples/spanner_rag_agent/README.md
+44 -4
View File
@@ -39,10 +39,10 @@ CREATE TABLE products (
productDescription STRING(MAX) NOT NULL,
productDescriptionEmbedding ARRAY<FLOAT32>,
createTime TIMESTAMP NOT NULL OPTIONS (
allow_commit_timestamp = true
),
inventoryCount INT64 NOT NULL,
priceInCents INT64,
allow_commit_timestamp = true
),
inventoryCount INT64 NOT NULL,
priceInCents INT64,
) PRIMARY KEY(categoryId, productId);
```
@@ -184,3 +184,43 @@ type.
* I'd like to buy a starter bike for my 3 year old child, can you show me the recommendation?
![Spanner RAG Sample Agent](Spanner_RAG_Sample_Agent.png)
## Which tool to use and When?
There are a few options to perform similarity search (see the `agent.py` for
implementation details):
1. Wraps the built-in `similarity_search` in the Spanner Toolset.
- 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.
2. Wraps the built-in `execute_sql` in the Spanner Toolset.
- `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.
3. Use the Spanner Toolset (and all the tools that come with it) directly.
- 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:
```py
toolset = SpannerToolset(
credentials_config=credentials_config,
tool_filter=["execute_sql", "get_table_schema"],
spanner_tool_settings=SpannerToolSettings(),
)
```
contributing/samples/spanner_rag_agent/agent.py
+74 -22
View File
@@ -22,10 +22,10 @@ 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.spanner_credentials import SpannerCredentialsConfig
from google.adk.tools.spanner.spanner_toolset import SpannerToolset
from google.adk.tools.tool_context import ToolContext
import google.auth
from google.auth.credentials import Credentials
@@ -67,15 +67,8 @@ else:
credentials=application_default_credentials
)
### Section 1: Get the built-in Spanner toolset ###
# Note that the built-in Spanner toolset is more flexible and generic. It is
# shown here for comparison and tutorial purposes.
spanner_toolset = SpannerToolset(
credentials_config=credentials_config, spanner_tool_settings=tool_settings
)
### Section 2: Extending the built-in Spanner Toolset for Custom Use Cases ###
### 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:
@@ -98,6 +91,10 @@ class SpannerRagSetting(BaseModel):
# 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] = [
@@ -115,10 +112,61 @@ class SpannerRagSetting(BaseModel):
RAG_SETTINGS = SpannerRagSetting()
# Create a wrapped function tool for the agent on top of the built-in Spanner
# toolset.
### (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
# embeded knowledge base stored in a Spanner database table.
# 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
@@ -134,10 +182,6 @@ def wrapped_spanner_execute_sql_tool(
Relevant product catalog content with sources
"""
# 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_query = f"""SELECT embeddings.values
FROM ML.PREDICT(
MODEL {RAG_SETTINGS.embedding_model_name},
@@ -188,7 +232,7 @@ def inspect_tool_params(
pass
### Section 3: Create the root agent ###
### Section 2: Create the root agent ###
root_agent = LlmAgent(
model="gemini-2.5-flash",
name="spanner_knowledge_base_agent",
@@ -197,19 +241,27 @@ root_agent = LlmAgent(
),
instruction="""
You are a helpful assistant that answers user questions about product-specific recommendations.
1. Always use the `wrapped_spanner_execute_sql_tool` tool to find relevant information.
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 Spanner toolset.
# # (Option 1)
# # Add customized Spanner tool based on the built-in similarity_search
# # in the Spanner toolset.
GoogleTool(
func=wrapped_spanner_execute_sql_tool,
func=wrapped_spanner_similarity_search,
credentials_config=credentials_config,
tool_settings=tool_settings,
),
# Add built-in Spanner toolset for comparison and tutorial purposes.
# spanner_toolset,
# # (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,
)