Skip to content

Add optional AlloyDB search support#110

Open
akshayjadiyanv wants to merge 1 commit into
NVIDIA-AI-Blueprints:mainfrom
akshayjadiyanv:alloydb-search-support
Open

Add optional AlloyDB search support#110
akshayjadiyanv wants to merge 1 commit into
NVIDIA-AI-Blueprints:mainfrom
akshayjadiyanv:alloydb-search-support

Conversation

@akshayjadiyanv

@akshayjadiyanv akshayjadiyanv commented May 20, 2026

Copy link
Copy Markdown

Summary

This PR adds optional AlloyDB AI vector search support for the Search Agent without changing the default Milvus-backed behavior.

  • Adds a new NAT function (alloydb_search) that embeds queries via NIM and runs cosine vector search against AlloyDB (pgvector + ScaNN).
  • Adds a dedicated config (configs/search-alloydb.yml) for AlloyDB deployments; the default configs/search.yml remains Milvus-backed and is unchanged.
  • Adds an AlloyDB seeder (scripts/seed_alloydb.py) mirroring the existing Milvus seeder.
  • Registers the new function in register.py and adds required dependencies (psycopg, pgvector, google-cloud-alloydb-connector).
  • Updates the agents Docker image and adds cloudbuild.yaml for native amd64 Cloud Build image builds.

Why

The blueprint currently supports Milvus for local/docker-compose workflows. This change adds an opt-in AlloyDB path for GCP deployments (for example Cloud Run + AlloyDB AI) while preserving existing Milvus defaults for adopters who do not use AlloyDB.

Usage

Default (unchanged):

nat serve --config_file configs/search.yml --port 8005

AlloyDB (opt-in):

python scripts/seed_alloydb.py
nat serve --config_file configs/search-alloydb.yml --port 8005

For Cloud Run, pass --config_file=configs/search-alloydb.yml and set AlloyDB env vars (ALLOYDB_USE_IAM_AUTH, ALLOYDB_INSTANCE_URI, etc.).

Test plan

  • Confirm configs/search.yml still uses milvus_retriever + nat_retriever (default behavior unchanged)
  • Run Search Agent locally with configs/search.yml against seeded Milvus
  • Seed AlloyDB with scripts/seed_alloydb.py and run Search Agent with configs/search-alloydb.yml
  • Verify Cloud Run deploy with configs/search-alloydb.yml and AlloyDB IAM auth (optional GCP validation)
  • Confirm recommendation agent and other configs still use Milvus unchanged

@akshayjadiyanv akshayjadiyanv marked this pull request as ready for review May 31, 2026 14:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant