diff --git a/spanner-adk-agents/README.md b/spanner-adk-agents/README.md new file mode 100644 index 0000000..854f5d3 --- /dev/null +++ b/spanner-adk-agents/README.md @@ -0,0 +1,96 @@ +# Build SQL and Semantic Search Agents with the ADK and Spanner + +This directory contains a sample demonstrating how to build a healthcare agent using the [Agent Development Kit (ADK)](https://adk.dev/) integration with **Cloud Spanner**. It is based on [this codelab](https://codelabs.developers.google.com/spanner-adk) (coming soon!). + +The agents in this sample show you how to: + +* Create a Spanner database and add an embedding column +* Configure ADK's SpannerToolset and SpannerToolSettings. +* Enable database metadata inspection, SQL execution, and vector similarity search tools. + +While this example shows a healthcare use case, the capabilities demonstrated in this sample can be applied to many industries. + +## Prerequisites + +1. A Google Cloud project with billing enabled. +2. The following APIs must be enabled: + ```bash + gcloud services enable spanner.googleapis.com aiplatform.googleapis.com + ``` + +## Environment variables + +To run the samples, set the following env vars: + +```shell +# Your Google Cloud Project ID +export GOOGLE_CLOUD_PROJECT=$(gcloud config get-value project) + +# Google Cloud location +export GOOGLE_CLOUD_LOCATION="us" + +# Instruct ADK to use Vertex AI rather than the public Gemini API +export GOOGLE_GENAI_USE_VERTEXAI=True + +# Your Spanner instance and database configuration +export SPANNER_INSTANCE_ID="healthcare" +export SPANNER_DATABASE_ID="medical-db" +``` + +## Create Spanner instance and database + +Run the following commands using the `gcloud` CLI to create a Spanner instance and DB to use for this sample: + +### Create the Spanner instance +gcloud spanner instances create $SPANNER_INSTANCE_ID \ + --config=regional-us-central1 \ + --description="ADK Sample Instance" \ + --edition=enterprise \ + --processing-units=1000 + +### Create the database +gcloud spanner databases create $SPANNER_DATABASE_ID \ + --instance=$SPANNER_INSTANCE_ID + + +## Setup + +1. Install Dependencies: +Clone the repository, navigate to this directory, and install the required Python packages: + +```shell +pip install -r requirements.txt +``` + +2. Create tables in your database and load data: + +First, execute the queries provided in `creaet_tables.sql` to create the `Providers`, `Patients`, `Appointments`, and `Prescriptions` tables and populate them with sample data. + +Next, run the queries in `embeddings.sql` to generate the vector embeddings and populate them into the `DoctorNotesEmbedding` column. + +**Important**: When running the `CREATE OR REPLACE MODEL TextEmbeddingModel` statement, ensure you replace `YOUR_PROJECT_ID` with your actual Google Cloud Project ID. + +## Running the Agents Locally + +This sample contains three different agents: + +* `basic_spanner_agent`: Answers metadata questions and executes SQL to summarize your database. +* `semantic_agent`: Performs vector similarity search on unstructured doctor notes. +* `secure_agent`: Demonstrates how to restrict agent access to only an allowed list of tables. + +To test any of the agents locally with the ADK Web UI, run the following command from the root of this sample directory: + +```shell +adk web --allow_origins="regex:.*" --session_service_uri="memory://" . +``` + +Navigate to http://127.0.0.1:8000 in your browser. Use the dropdown at the top of the interface to switch between the different agents and interact with them. + +## Cleanup + +To avoid incurring unexpected charges to your Google Cloud billing account, make sure to delete the Spanner instance when you are done testing this sample. + +Deleting the instance will also automatically delete the `medical-db` database and all of its data. + +```bash +gcloud spanner instances delete $SPANNER_INSTANCE_ID --quiet diff --git a/spanner-adk-agents/basic_spanner_agent/agent.py b/spanner-adk-agents/basic_spanner_agent/agent.py new file mode 100644 index 0000000..02e333f --- /dev/null +++ b/spanner-adk-agents/basic_spanner_agent/agent.py @@ -0,0 +1,85 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +import google.auth + +from google.auth.transport.requests import Request +from google.adk.agents import Agent +from google.adk.tools.spanner import client as spanner_client_module +from google.adk.tools.spanner.settings import SpannerToolSettings, Capabilities +from google.adk.tools.spanner.spanner_credentials import SpannerCredentialsConfig +from google.adk.tools.spanner.spanner_toolset import SpannerToolset + +import os +import sys + +# --- Environment Variable Validation --- +REQUIRED_ENV_VARS = [ + "GOOGLE_CLOUD_PROJECT", + "GOOGLE_CLOUD_LOCATION", + "GOOGLE_GENAI_USE_VERTEXAI", + "SPANNER_INSTANCE_ID", + "SPANNER_DATABASE_ID" +] + +missing_vars = [var for var in REQUIRED_ENV_VARS if not os.environ.get(var)] +if missing_vars: + print(f"\n[ERROR] Missing required environment variables: {', '.join(missing_vars)}", flush=True) + print("Please set them before running the agent. See the README.md for instructions.\n", flush=True) + sys.exit(1) + +PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT") +INSTANCE_ID = os.environ.get("SPANNER_INSTANCE_ID", "healthcare") +DATABASE_ID = os.environ.get("SPANNER_DATABASE_ID", "medical-db") + +# --- Cloud Shell Workaround --- +# If you are running this in Cloud Shell, you need to apply a small patch to prevent +# the local proxy from deadlocking Spanner's gRPC cleanup. +# You do not need this in standard production environments like Cloud Run or your local laptop. +if hasattr(spanner_client_module, "_close_spanner_resources"): + spanner_client_module._close_spanner_resources = lambda *args, **kwargs: None + +# --- Auth --- +try: + application_default_credentials, _ = google.auth.default() + if not application_default_credentials.valid: + application_default_credentials.refresh(Request()) +except Exception as e: + print(f"\n[ERROR] Failed to authenticate: {e}", flush=True) + sys.exit(1) + +# --- Spanner Tool config --- +credentials_config = SpannerCredentialsConfig(credentials=application_default_credentials) +tool_settings = SpannerToolSettings(capabilities=[Capabilities.DATA_READ]) +spanner_toolset = SpannerToolset(credentials_config=credentials_config, spanner_tool_settings=tool_settings) + +root_agent = Agent( + model="gemini-3.8-flash", + name="spanner_healthcare_agent", + description="Agent to answer questions about Spanner database and execute SQL queries.", + instruction=f""" + You are a data assistant agent with access to several Spanner tools. + Make use of those tools to answer the user's questions. + + When using your tools, always use the following default database configuration: + - project_id: {PROJECT_ID} + - instance_id: {INSTANCE_ID} + - database_id: {DATABASE_ID} + """, + tools=[ + spanner_toolset, + ], +) \ No newline at end of file diff --git a/spanner-adk-agents/create_tables.sql b/spanner-adk-agents/create_tables.sql new file mode 100644 index 0000000..b5db19e --- /dev/null +++ b/spanner-adk-agents/create_tables.sql @@ -0,0 +1,104 @@ +-- Copyright 2026 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +-- 1. Create the independent Providers table +CREATE TABLE Providers ( + ProviderId INT64 NOT NULL, + ProviderName STRING(MAX), + Specialty STRING(MAX) +) PRIMARY KEY (ProviderId); + +-- 2. Create the parent Patients table +CREATE TABLE Patients ( + PatientId INT64 NOT NULL, + FirstName STRING(MAX), + LastName STRING(MAX), + DateOfBirth DATE +) PRIMARY KEY (PatientId); + +-- 3. Create the interleaved Appointments table +CREATE TABLE Appointments ( + PatientId INT64 NOT NULL, + AppointmentId INT64 NOT NULL, + ProviderId INT64, + AppointmentDate DATE, + Status STRING(MAX), + DoctorNotes STRING(MAX), + DoctorNotesEmbedding ARRAY(vector_length=>768) +) PRIMARY KEY (PatientId, AppointmentId), + INTERLEAVE IN PARENT Patients ON DELETE CASCADE; + +-- 4. Create the interleaved Prescriptions table +CREATE TABLE Prescriptions ( + PatientId INT64 NOT NULL, + PrescriptionId INT64 NOT NULL, + ProviderId INT64 NOT NULL, + MedicationName STRING(MAX), + DatePrescribed DATE +) PRIMARY KEY (PatientId, PrescriptionId), + INTERLEAVE IN PARENT Patients ON DELETE CASCADE; + +-- 5. Insert data into each table +INSERT INTO Providers (ProviderId, ProviderName, Specialty) +VALUES + (1, 'Dr. Aris Thorne', 'Culinary Diagnostics'), + (2, 'Dr. Beatrice Plum', 'Chronological Confusion'), + (3, 'Dr. Caspian Vane', 'Gravity Resistance'), + (4, 'Dr. Delilah Frost', 'Acute Dessert Therapy'), + (5, 'Dr. Eldon Rook', 'Somnambulant Gymnastics'), + (6, 'Dr. Fiona Gale', 'Over-enthusiastic Sneezing'), + (7, 'Dr. Gideon Vance', 'Extreme Couch Potato-ism'), + (8, 'Dr. Hazel Quinn', 'Spontaneous Melody Outbreaks'), + (9, 'Dr. Ignatius Flint', 'Stubbed Toe Trauma'), + (10, 'Dr. Juniper Slate', 'Advanced Broccoli Administration'); + +INSERT INTO Patients (PatientId, FirstName, LastName, DateOfBirth) +VALUES + (1, 'Barnaby', 'Quigley', DATE '1982-04-12'), + (2, 'Seraphina', 'Pockets', DATE '1995-11-23'), + (3, 'Thaddeus', 'Plumbob', DATE '1978-01-30'), + (4, 'Marigold', 'Swoon', DATE '2002-08-14'), + (5, 'Silas', 'Fiddlewood', DATE '1965-06-05'), + (6, 'Clementine', 'Fizz', DATE '1988-12-10'), + (7, 'Orville', 'Snipe', DATE '1971-03-22'), + (8, 'Rosalind', 'Furlong', DATE '1999-07-07'), + (9, 'Percival', 'Gout', DATE '1955-09-18'), + (10, 'Elara', 'Moonbeam', DATE '2010-02-28'); + +INSERT INTO Appointments (PatientId, AppointmentId, ProviderId, AppointmentDate, Status, DoctorNotes) +VALUES + (1, 101, 1, DATE '2026-10-01', 'Completed', 'Patient complains of tasting the color blue. Prescribed 14 hours of video games and a large pizza.'), + (2, 102, 3, DATE '2026-10-02', 'Completed', 'Patient accidentally swallowed a cloud. Floating slightly above the exam table. Needs a heavy lunch to weigh her down.'), + (3, 103, 2, DATE '2026-10-03', 'Completed', 'Patient has developed a severe allergy to Mondays. Breaking out in hives when looking at a calendar.'), + (4, 104, 6, DATE '2026-10-04', 'Completed', 'Patient reports excessive glitter in bloodstream after crafting accident. Sparkles violently when sneezing.'), + (5, 105, 5, DATE '2026-10-05', 'Completed', 'Diagnosed with resting confused face. Patient forgot why he came to the clinic in the first place.'), + (6, 106, 4, DATE '2026-10-06', 'Scheduled', 'Severe case of ice cream withdrawal. Symptoms include whining, shivering, and aggressively pointing at freezers.'), + (7, 107, 7, DATE '2026-10-07', 'Completed', 'Patient left leg has fallen asleep and refuses to wake up without a bedtime story.'), + (8, 108, 8, DATE '2026-10-08', 'Completed', 'Uncontrollable urge to tap dance when hearing elevator music. Ankles are showing signs of extreme wear.'), + (9, 109, 9, DATE '2026-10-09', 'Scheduled', 'Patient believes his eyebrows are trying to escape. Taped them down pending further review.'), + (10, 110, 10, DATE '2026-10-10', 'Completed', 'Diagnosed with acute vegetable aversion. Emits a high-pitched frequency when placed within 10 feet of broccoli.'); + +INSERT INTO Prescriptions (PatientId, PrescriptionId, ProviderId, MedicationName, DatePrescribed) +VALUES + (1, 201, 1, 'Extra Cheese Pepperoni Pizza (Taken orally)', DATE '2026-10-01'), + (2, 202, 3, 'Lead-weighted Boots (Wear daily)', DATE '2026-10-02'), + (3, 203, 2, 'Time Machine set to Tuesday (Use once)', DATE '2026-10-03'), + (4, 204, 6, 'Vacuum Cleaner on Reverse Mode (Apply to nose)', DATE '2026-10-04'), + (5, 205, 5, 'A Map and a Compass (Consult twice daily)', DATE '2026-10-05'), + (6, 206, 4, 'Three Scoops of Neapolitan Ice Cream (Stat!)', DATE '2026-10-06'), + (7, 207, 7, 'Collection of Fairy Tales (Read to leg at 9PM)', DATE '2026-10-07'), + (8, 208, 8, 'Noise Cancelling Headphones (Wear near lobbies)', DATE '2026-10-08'), + (9, 209, 9, 'Heavy Duty Masking Tape (Apply to forehead)', DATE '2026-10-09'), + (10, 210, 10, 'Chocolate-Covered Broccoli (To trick the system, eat with caution)', DATE '2026-10-10'); + diff --git a/spanner-adk-agents/embeddings.sql b/spanner-adk-agents/embeddings.sql new file mode 100644 index 0000000..64702c7 --- /dev/null +++ b/spanner-adk-agents/embeddings.sql @@ -0,0 +1,34 @@ +-- Copyright 2026 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +-- IMPORTANT: Before copying this into Spanner Studio, replace YOUR_PROJECT_ID with the ID of the Google Cloud project you are using to run this sample + +-- 1. Create a text embedding model +CREATE OR REPLACE MODEL TextEmbeddingModel +INPUT(content STRING(MAX)) +OUTPUT(embeddings STRUCT>) +REMOTE OPTIONS( + endpoint = '//aiplatform.googleapis.com/projects/YOUR_PROJECT_ID/locations/us-central1/publishers/google/models/text-embedding-005' +); + +-- 2. Generate and insert vector embeddings +UPDATE Appointments +SET DoctorNotesEmbedding = ( + SELECT embeddings.values + FROM ML.PREDICT( + MODEL TextEmbeddingModel, + (SELECT DoctorNotes AS content) + ) +) +WHERE DoctorNotes IS NOT NULL AND DoctorNotesEmbedding IS NULL; \ No newline at end of file diff --git a/spanner-adk-agents/requirements.txt b/spanner-adk-agents/requirements.txt new file mode 100644 index 0000000..b7b7483 --- /dev/null +++ b/spanner-adk-agents/requirements.txt @@ -0,0 +1,4 @@ +google-cloud-spanner>=3.64.0 +google-adk[spanner] +google-genai +google-auth \ No newline at end of file diff --git a/spanner-adk-agents/secure_agent/agent.py b/spanner-adk-agents/secure_agent/agent.py new file mode 100644 index 0000000..5c0298f --- /dev/null +++ b/spanner-adk-agents/secure_agent/agent.py @@ -0,0 +1,119 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +import google.auth + +from google.auth.credentials import Credentials +from google.auth.transport.requests import Request +from google.adk.agents import Agent +from google.adk.tools import ToolContext +from google.adk.tools.spanner import query_tool +from google.adk.tools.google_tool import GoogleTool +from google.adk.tools.spanner import client as spanner_client_module +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 + +# --- Environment Variable Validation --- +REQUIRED_ENV_VARS = [ + "GOOGLE_CLOUD_PROJECT", + "GOOGLE_CLOUD_LOCATION", + "GOOGLE_GENAI_USE_VERTEXAI", + "SPANNER_INSTANCE_ID", + "SPANNER_DATABASE_ID" +] + +missing_vars = [var for var in REQUIRED_ENV_VARS if not os.environ.get(var)] +if missing_vars: + print(f"\n[ERROR] Missing required environment variables: {', '.join(missing_vars)}", flush=True) + print("Please set them before running the agent. See the README.md for instructions.\n", flush=True) + sys.exit(1) + +PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT") +INSTANCE_ID = os.environ.get("SPANNER_INSTANCE_ID", "healthcare") +DATABASE_ID = os.environ.get("SPANNER_DATABASE_ID", "medical-db") + +# --- Cloud Shell Workaround --- +# If you are running this in Cloud Shell, you need to apply a small patch to prevent +# the local proxy from deadlocking Spanner's gRPC cleanup. +# You do not need this in standard production environments like Cloud Run or your local laptop. +if hasattr(spanner_client_module, "_close_spanner_resources"): + spanner_client_module._close_spanner_resources = lambda *args, **kwargs: None + +# --- Auth --- +try: + application_default_credentials, _ = google.auth.default() + if not application_default_credentials.valid: + application_default_credentials.refresh(Request()) +except Exception as e: + print(f"\n[ERROR] Failed to authenticate: {e}", flush=True) + sys.exit(1) + +# --- Spanner Tool config --- +credentials_config = SpannerCredentialsConfig( + credentials=application_default_credentials +) + +# Revoke general SQL access by passing an empty capabilities list +secure_tool_settings = SpannerToolSettings(capabilities=[]) + +secure_spanner_toolset = SpannerToolset( + credentials_config=credentials_config, + spanner_tool_settings=secure_tool_settings +) + +# Define a strict allowlist of tables the agent is permitted to access +ALLOWED_TABLES = {"Appointments", "Prescriptions"} + +async def secure_count_rows_tool( + table_name: str, + credentials: Credentials, + settings: SpannerToolSettings, + tool_context: ToolContext, +): + """Counts the total number of rows for a specified table.""" + + # Guard against SQL injection and unauthorized access + if table_name not in ALLOWED_TABLES: + return {"status": "ERROR", "message": f"Unauthorized table: {table_name}"} + + return await query_tool.execute_sql( + project_id=PROJECT_ID, + instance_id=INSTANCE_ID, + database_id=DATABASE_ID, + query=f"SELECT count(*) FROM {table_name}", + credentials=credentials, + settings=settings, + tool_context=tool_context, + ) + +# Give the agent our secured tool +root_agent = Agent( + model="gemini-3.8-flash", + name="secure_healthcare_agent", + description="Agent to count rows in allowed Spanner tables.", + instruction=""" + You are a database agent with access to a custom SQL tool that can run on specified tables. + """, + tools=[ + secure_spanner_toolset, + GoogleTool( + func=secure_count_rows_tool, + credentials_config=credentials_config, + tool_settings=secure_tool_settings, + ), + ], +) \ No newline at end of file diff --git a/spanner-adk-agents/semantic_agent/agent.py b/spanner-adk-agents/semantic_agent/agent.py new file mode 100644 index 0000000..5ba29ad --- /dev/null +++ b/spanner-adk-agents/semantic_agent/agent.py @@ -0,0 +1,121 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +import asyncio +import google.auth + +from google.auth.transport.requests import Request +from google.adk.agents import Agent +from google.adk.tools.spanner import utils as spanner_utils +from google.adk.tools.spanner import client as spanner_client_module +from google.adk.tools.spanner.settings import SpannerToolSettings, Capabilities, SpannerVectorStoreSettings +from google.adk.tools.spanner.spanner_credentials import SpannerCredentialsConfig +from google.adk.tools.spanner.spanner_toolset import SpannerToolset +from google.genai import Client + +# --- Environment Variable Validation --- +REQUIRED_ENV_VARS = [ + "GOOGLE_CLOUD_PROJECT", + "GOOGLE_CLOUD_LOCATION", + "GOOGLE_GENAI_USE_VERTEXAI", + "SPANNER_INSTANCE_ID", + "SPANNER_DATABASE_ID" +] + +missing_vars = [var for var in REQUIRED_ENV_VARS if not os.environ.get(var)] +if missing_vars: + print(f"\n[ERROR] Missing required environment variables: {', '.join(missing_vars)}", flush=True) + print("Please set them before running the agent. See the README.md for instructions.\n", flush=True) + sys.exit(1) + +PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT") +INSTANCE_ID = os.environ.get("SPANNER_INSTANCE_ID", "healthcare") +DATABASE_ID = os.environ.get("SPANNER_DATABASE_ID", "medical-db") + +# --- Cloud Shell Workaround --- +# If you are running this in Cloud Shell, you need to apply a small patch to prevent +# the local proxy from deadlocking Spanner's gRPC cleanup. +# You do not need this in standard production environments like Cloud Run or your local laptop. +if hasattr(spanner_client_module, "_close_spanner_resources"): + spanner_client_module._close_spanner_resources = lambda *args, **kwargs: None + +# --- Auth and credentials config --- +try: + application_default_credentials, _ = google.auth.default() + if not application_default_credentials.valid: + application_default_credentials.refresh(Request()) +except Exception as e: + print(f"\n[ERROR] Failed to authenticate: {e}", flush=True) + sys.exit(1) + +credentials_config = SpannerCredentialsConfig( + credentials=application_default_credentials +) + +# Because text-embedding-005 uses a regional endpoint (i.e. 'us-central1') +# and gemini-3.8-flash uses a multi-regional endpoint (i.e. 'us') +# We need the following patch to pass the correct region to the embedding model +regional_client = Client(vertexai=True, project=PROJECT_ID, location="us-central1") +async def _patched_embed_contents_async( + vertex_ai_embedding_model_name: str, + contents: list[str], + output_dimensionality: int | None = None, + genai_client: Client | None = None + ): + return await asyncio.to_thread( + spanner_utils.embed_contents, vertex_ai_embedding_model_name, contents, + output_dimensionality=output_dimensionality, genai_client=regional_client, + ) +spanner_utils.embed_contents_async = _patched_embed_contents_async + +my_vector_store_settings = SpannerVectorStoreSettings( + project_id=PROJECT_ID, + instance_id=INSTANCE_ID, + database_id=DATABASE_ID, + table_name="Appointments", + content_column="DoctorNotes", + embedding_column="DoctorNotesEmbedding", + vector_length=768, + vertex_ai_embedding_model_name="text-embedding-005", + selected_columns=["PatientId", "ProviderId", "AppointmentDate", "DoctorNotes"], + nearest_neighbors_algorithm="EXACT_NEAREST_NEIGHBORS", + top_k=3, + distance_type="COSINE", +) + +my_tool_settings = SpannerToolSettings( + capabilities=[Capabilities.DATA_READ], + vector_store_settings=my_vector_store_settings, +) + +my_spanner_toolset = SpannerToolset( + credentials_config=credentials_config, + spanner_tool_settings=my_tool_settings, + tool_filter=["vector_store_similarity_search"], +) + +root_agent = Agent( + model="gemini-3.8-flash", + name="doctor_notes_agent", + description="Semantic search agent for medical notes.", + instruction=""" + You are a medical assistant that returns patient and provider IDs based on the query. + 1. Always use the `spanner_vector_store_similarity_search` tool to find associated IDs. + 2. If no relevant information is found, state that no patient notes matched the request. + 3. Present the patient ID, provider ID, and relevant doctor notes to make it clear why each patient and provider were returned for a particular search. + """, + tools=[my_spanner_toolset], +) \ No newline at end of file