feat: add Vertex AI ADC embedding provider (google-vertex recipe)#729
Closed
lucha0404 wants to merge 2 commits into
Closed
feat: add Vertex AI ADC embedding provider (google-vertex recipe)#729lucha0404 wants to merge 2 commits into
lucha0404 wants to merge 2 commits into
Conversation
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.
8 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
The existing
googlerecipe usesGOOGLE_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-vertexrecipe that authenticates through Vertex AI's standard ADC chain (via google-auth-library), supporting:gcloud auth application-default loginGOOGLE_APPLICATION_CREDENTIALSenvUse 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.src/core/ai/types.tsImplementationenum with'native-google-vertex'src/core/ai/recipes/google-vertex.ts(new, 36 lines)gemini-embedding-001+text-multilingual-embedding-002+text-embedding-005, matryoshka 1536d defaultsrc/core/ai/recipes/index.tsgoogleVertexinALL[]src/core/ai/gateway.tsinstantiateEmbeddingswitch case +createVerteximportsrc/core/ai/dims.tsproviderOptionswithvertex/google/googleVertexnamespaces (v4 + v5 SDK compat)Dependency:
@ai-sdk/google-vertex@^4.0.0(matches@ai-sdk/google@^3.0.xpeer 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.tsinsertChunkshad 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 sameawait import('./ai/gateway.ts')patterninitSchemaalready uses (L153). Pre-existing bug, surfaced when running gbrain end-to-end on a non-OpenAI provider.Usage
After this lands:
Test plan
bun run typecheckclean (no errors)gbrain doctorembedding_providerprobe returns 1536d, dim-aligned with schemaput_page→searchon Supabase Postgres, cross-machine consistent (write from Windows + read from Linux GCE, same RRF score 0.4489)@ai-sdk/google-vertex@4.0.122(current 4.x stable)Compatibility notes
dims.tsincludes threeproviderOptionsnamespaces (vertex,google,googleVertex) because:@ai-sdk/google-vertex@4.xparsesvertexfirst, then falls back togoogle@ai-sdk/google-vertex@5.xcanary parsesgoogleVertexfirst, thenvertexPassing 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.
Need help on this PR? Tag
@codesmithwith what you need.