Skip to content

feat: add Vertex AI ADC embedding provider (google-vertex recipe)#729

Closed
lucha0404 wants to merge 2 commits into
garrytan:masterfrom
lucha0404:vertex-ai-adc
Closed

feat: add Vertex AI ADC embedding provider (google-vertex recipe)#729
lucha0404 wants to merge 2 commits into
garrytan:masterfrom
lucha0404:vertex-ai-adc

Conversation

@lucha0404

@lucha0404 lucha0404 commented May 8, 2026

Copy link
Copy Markdown

Why

The existing google recipe uses GOOGLE_GENERATIVE_AI_API_KEY (AI Studio public API key). This works for personal use but doesn't fit production GCP deployments that already have a service account / GCE metadata-based ADC chain.

This PR adds a parallel google-vertex recipe that authenticates through Vertex AI's standard ADC chain (via google-auth-library), supporting:

  • GCE / GKE: automatic credentials via metadata service
  • Local dev: gcloud auth application-default login
  • Production: standard GOOGLE_APPLICATION_CREDENTIALS env

Use case: GCP-deployed gbrain instances (e.g. on a Compute Engine VM) can drop the AI Studio API key entirely and reuse the project's service account.

What

Five-file extension to the existing pluggable provider framework. No changes to embedding.ts (still a thin wrapper); all wiring goes through the documented "new native provider = recipe + factory in gateway.ts" pattern.

File Change
src/core/ai/types.ts extend Implementation enum with 'native-google-vertex'
src/core/ai/recipes/google-vertex.ts (new, 36 lines) recipe declares gemini-embedding-001 + text-multilingual-embedding-002 + text-embedding-005, matryoshka 1536d default
src/core/ai/recipes/index.ts register googleVertex in ALL[]
src/core/ai/gateway.ts instantiateEmbedding switch case + createVertex import
src/core/ai/dims.ts providerOptions with vertex / google / googleVertex namespaces (v4 + v5 SDK compat)

Dependency: @ai-sdk/google-vertex@^4.0.0 (matches @ai-sdk/google@^3.0.x peer family / ai@^6.0.x).

Total: 132 insertions, 1 deletion across 8 files (incl. package.json + bun.lock).

Bonus fix (last commit, fully independent — feel free to cherry-pick separately)

src/core/postgres-engine.ts insertChunks had a hardcoded 'text-embedding-3-large' fallback regardless of the active provider, so chunks recorded OpenAI metadata even when Vertex / Voyage / etc. did the embed. The fix mirrors the same await import('./ai/gateway.ts') pattern initSchema already uses (L153). Pre-existing bug, surfaced when running gbrain end-to-end on a non-OpenAI provider.

Usage

After this lands:

# config.json
{
  "engine": "postgres",
  "database_url": "postgresql://...",
  "embedding_model": "google-vertex:gemini-embedding-001"
}

# env (or systemd Environment=)
GBRAIN_VERTEX_PROJECT=your-gcp-project
GBRAIN_VERTEX_LOCATION=us-central1   # optional, defaults to us-central1

gbrain init
gbrain doctor   # → google-vertex:gemini-embedding-001 ✓ NNNNms, 1536 dims, DB aligned

Test plan

  • bun run typecheck clean (no errors)
  • gbrain doctor embedding_provider probe returns 1536d, dim-aligned with schema
  • ✅ End-to-end put_pagesearch on Supabase Postgres, cross-machine consistent (write from Windows + read from Linux GCE, same RRF score 0.4489)
  • ✅ Tested with @ai-sdk/google-vertex@4.0.122 (current 4.x stable)

Compatibility notes

dims.ts includes three providerOptions namespaces (vertex, google, googleVertex) because:

  • @ai-sdk/google-vertex@4.x parses vertex first, then falls back to google
  • @ai-sdk/google-vertex@5.x canary parses googleVertex first, then vertex

Passing all three keeps the dim-passthrough working across both major versions, at zero runtime cost (parseProviderOptions short-circuits on first match).

Rebased onto v0.28.12

Branch is rebased onto upstream/master HEAD (bca993e v0.28.12) so v0.28.11 multimodal embedding routing + v0.28.12 LongMemEval harness are preserved alongside the new Vertex AI provider.


View in Codesmith
Need help on this PR? Tag @codesmith with what you need.

  • Let Codesmith autofix CI failures and bot reviews

lucha0404 added 2 commits May 8, 2026 13:10
Adds 'native-google-vertex' Implementation alongside existing 'native-google'
(GOOGLE_GENERATIVE_AI_API_KEY / AI Studio). Vertex AI authenticates through
google-auth-library's full ADC chain:
  - GCE / GKE: metadata service (zero config)
  - Local: gcloud auth application-default login
  - Production: GOOGLE_APPLICATION_CREDENTIALS env

Use case: GCP-deployed gbrain instances reuse their service account instead
of provisioning a separate AI Studio API key.

Patch surface (5 files):
  - types.ts: extend Implementation enum
  - recipes/google-vertex.ts: new recipe (gemini-embedding-001, matryoshka 1536d)
  - recipes/index.ts: register googleVertex in ALL[]
  - gateway.ts: instantiateEmbedding switch case + createVertex import
  - dims.ts: outputDimensionality passthrough via 'vertex'/'google' namespaces

Dependency: @ai-sdk/google-vertex@^4.0.0 (matches ai@^6 + @ai-sdk/google@^3 family)

Usage:
  ~/.gbrain/config.json:
    "embedding_model": "google-vertex:gemini-embedding-001"
  env:
    GBRAIN_VERTEX_PROJECT=<gcp-project>
    GBRAIN_VERTEX_LOCATION=us-central1  (optional, default us-central1)

Verified end-to-end on Supabase Postgres + Vertex AI Gemini 1536d via
'gbrain doctor': google-vertex:gemini-embedding-001 ✓ ~1500ms, 1536 dims,
DB aligned.
…openai

postgres-engine.ts insertChunks fallback was hardcoded 'text-embedding-3-large'
so any chunk without explicit chunk.model recorded under the OpenAI name
regardless of which provider actually generated the embedding. Pre-existing
bug, surfaces visibly when running gbrain end-to-end on a non-OpenAI provider.

Mirrors the same dynamic-import pattern initSchema (L153) already uses.

Independent of the Vertex AI patch — feel free to cherry-pick separately.
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