From 9845460a1d1feeb0f90ebf4c54f9b6f1f78742f2 Mon Sep 17 00:00:00 2001 From: duwenxin Date: Sun, 1 Mar 2026 20:11:43 -0500 Subject: [PATCH 1/7] feat(embeddingModel): Add Backend API selection fields --- .../configuration/embedding-models/_index.md | 2 + .../configuration/embedding-models/gemini.md | 70 +++++++++++++----- internal/embeddingmodels/gemini/gemini.go | 72 +++++++++++++++---- 3 files changed, 113 insertions(+), 31 deletions(-) diff --git a/docs/en/documentation/configuration/embedding-models/_index.md b/docs/en/documentation/configuration/embedding-models/_index.md index b5bb7e985be2..40cf8fedec7d 100644 --- a/docs/en/documentation/configuration/embedding-models/_index.md +++ b/docs/en/documentation/configuration/embedding-models/_index.md @@ -73,6 +73,7 @@ kind: tool name: insert_embedding type: postgres-sql source: my-pg-instance +description: Insert a new document into the database. statement: | INSERT INTO documents (content, embedding) VALUES ($1, $2); @@ -92,6 +93,7 @@ kind: tool name: search_embedding type: postgres-sql source: my-pg-instance +description: Search for documents in the database. statement: | SELECT id, content, embedding <-> $1 AS distance FROM documents diff --git a/docs/en/documentation/configuration/embedding-models/gemini.md b/docs/en/documentation/configuration/embedding-models/gemini.md index 92cb53589ada..b9d481b25a51 100644 --- a/docs/en/documentation/configuration/embedding-models/gemini.md +++ b/docs/en/documentation/configuration/embedding-models/gemini.md @@ -3,7 +3,8 @@ title: "Gemini Embedding" type: docs weight: 1 description: > - Use Google's Gemini models to generate high-performance text embeddings for vector databases. + Use Google's Gemini models to generate high-performance text embeddings for + vector databases. --- ## About @@ -13,15 +14,18 @@ high-dimensional vectors. ### Authentication -Toolbox uses your [Application Default Credentials -(ADC)][adc] to authorize with the -Gemini API client. +Toolbox supports two authentication modes: -Optionally, you can use an [API key][api-key] obtain an API -Key from the [Google AI Studio][ai-studio]. +1. **Google AI (API Key):** Used if `useVertexAI` is `false` or unset. You must + provide `apiKey` (or set `GOOGLE_API_KEY`/`GEMINI_API_KEY` environment + variables). This uses the [Google AI Studio][ai-studio] backend. +2. **Vertex AI (ADC):** Used if `useVertexAI` is `true`. This uses [Application + Default Credentials (ADC)][adc]. When using this mode, you **must** specify + `project` and `location` (or set + `GOOGLE_CLOUD_PROJECT`/`GOOGLE_CLOUD_LOCATION` environment variables). -We recommend using an API key for testing and using application default -credentials for production. +We recommend using an API key for quick testing and using Vertex AI with ADC for +production environments. [adc]: https://cloud.google.com/docs/authentication#adc [api-key]: https://ai.google.dev/gemini-api/docs/api-key#api-keys @@ -41,14 +45,20 @@ to your database source. The `dimension` field must match the expected size of your database column (e.g., a `vector(768)` column in PostgreSQL). This setting is supported by newer models since 2024 only. You cannot set this value if using the earlier model -(`models/embedding-001`). Check out [available Gemini models][modellist] for more -information. +(`models/embedding-001`). Check out [available Gemini models][modellist] for +more information. [modellist]: - https://docs.cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings#supported-models + https://docs.cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings#supported-models + ## Example +### Using Google AI + +Google AI uses API Key for authentication. You can get an API key from [Google +AI Studio][ai-studio]. + ```yaml kind: embeddingModel name: gemini-model @@ -58,16 +68,38 @@ apiKey: ${GOOGLE_API_KEY} dimension: 768 ``` +### Using Vertex AI + +Vertex AI uses Application Default Credentials (ADC) for authentication. Learn +how to set up ADC [here][adc]. + +```yaml +kind: embeddingModels +name: gemini-model +type: gemini +model: text-embedding-004 +useVertexAI: true +project: ${GOOGLE_CLOUD_PROJECT} +location: us-central1 +dimension: 768 +``` + +[adc]: https://docs.cloud.google.com/docs/authentication/provide-credentials-adc + + {{< notice tip >}} -Use environment variable replacement with the format ${ENV_NAME} -instead of hardcoding your secrets into the configuration file. + Use environment variable replacement with the format + ${ENV_NAME} instead of hardcoding your secrets into the configuration file. {{< /notice >}} ## Reference -| **field** | **type** | **required** | **description** | -|-----------|:--------:|:------------:|--------------------------------------------------------------| -| type | string | true | Must be `gemini`. | -| model | string | true | The Gemini model ID to use (e.g., `gemini-embedding-001`). | -| apiKey | string | false | Your API Key from Google AI Studio. | -| dimension | integer | false | The number of dimensions in the output vector (e.g., `768`). | +| **field** | **type** | **required** | **description** | +| ----------- | :------: | :----------: | ------------------------------------------------------------- | +| type | string | true | Must be `gemini`. | +| model | string | true | The Gemini model ID to use (e.g., `text-embedding-004`). | +| useVertexAI | boolean | true | Set to `true` to use Vertex AI. Default is false (Google AI). | +| project | string | false | GCP Project ID (required if `useVertexAI` is `true`). | +| location | string | false | GCP Location (required if `useVertexAI` is `true`). | +| apiKey | string | false | Your API Key from Google AI Studio. | +| dimension | integer | false | The number of dimensions in the output vector (e.g., `768`). | diff --git a/internal/embeddingmodels/gemini/gemini.go b/internal/embeddingmodels/gemini/gemini.go index 2d0162e8f519..bd578b4946e3 100644 --- a/internal/embeddingmodels/gemini/gemini.go +++ b/internal/embeddingmodels/gemini/gemini.go @@ -18,6 +18,7 @@ import ( "context" "fmt" "net/http" + "os" "github.com/googleapis/genai-toolbox/internal/embeddingmodels" "github.com/googleapis/genai-toolbox/internal/util" @@ -30,11 +31,14 @@ const EmbeddingModelType string = "gemini" var _ embeddingmodels.EmbeddingModelConfig = Config{} type Config struct { - Name string `yaml:"name" validate:"required"` - Type string `yaml:"type" validate:"required"` - Model string `yaml:"model" validate:"required"` - ApiKey string `yaml:"apiKey"` - Dimension int32 `yaml:"dimension"` + Name string `yaml:"name" validate:"required"` + Type string `yaml:"type" validate:"required"` + Model string `yaml:"model" validate:"required"` + UseVertexAI bool `yaml:"useVertexAI" validate:"required"` + ApiKey string `yaml:"apiKey"` + Project string `yaml:"project"` + Location string `yaml:"location"` + Dimension int32 `yaml:"dimension"` } // Returns the embedding model type @@ -44,10 +48,55 @@ func (cfg Config) EmbeddingModelConfigType() string { // Initialize a Gemini embedding model func (cfg Config) Initialize(ctx context.Context) (embeddingmodels.EmbeddingModel, error) { - // Get client configs configs := &genai.ClientConfig{} - if cfg.ApiKey != "" { - configs.APIKey = cfg.ApiKey + + // Get API Key + apiKey := cfg.ApiKey + if apiKey == "" { + apiKey = os.Getenv("GOOGLE_API_KEY") + if apiKey == "" { + apiKey = os.Getenv("GEMINI_API_KEY") + } + } + + // Determine the Backend + if cfg.UseVertexAI { + // VertexAI API uses ADC for authentication. + // ADC requires `Project` and `Location` to be set. + configs.Backend = genai.BackendVertexAI + + configs.Project = cfg.Project + if configs.Project == "" { + configs.Project = os.Getenv("GOOGLE_CLOUD_PROJECT") + } + + configs.Location = cfg.Location + if configs.Location == "" { + configs.Location = os.Getenv("GOOGLE_CLOUD_LOCATION") + if configs.Location == "" { + configs.Location = "us-central1" + } + } + + if configs.Project == "" { + return nil, fmt.Errorf("vertex AI mode forced, but no project ID found in YAML or GOOGLE_CLOUD_PROJECT") + } + + if configs.Location == "" { + return nil, fmt.Errorf("vertex AI mode forced, but no location found in YAML or GOOGLE_CLOUD_LOCATION") + } + + } else if apiKey != "" { + // Using Gemini API, which uses API Key for authentication. + configs.Backend = genai.BackendGeminiAPI + configs.APIKey = apiKey + + } else { + // Missing both credentials + return nil, fmt.Errorf("missing credentials for Gemini embedding: " + + "For Google AI: Provide 'apiKey' in YAML or set GOOGLE_API_KEY/GEMINI_API_KEY env vars. " + + "For Vertex AI: Set 'useVertexAI: true' and provide 'project'/'location' in YAML or via GOOGLE_CLOUD_PROJECT/GOOGLE_CLOUD_LOCATION env vars. " + + "See documentation for details: https://googleapis.github.io/genai-toolbox/resources/embeddingmodels/gemini/") } ua, err := util.UserAgentFromContext(ctx) @@ -63,14 +112,13 @@ func (cfg Config) Initialize(ctx context.Context) (embeddingmodels.EmbeddingMode // Create new Gemini API client client, err := genai.NewClient(ctx, configs) if err != nil { - return nil, fmt.Errorf("unable to create Gemini API client") + return nil, fmt.Errorf("unable to create Gemini API client: %w", err) } - m := &EmbeddingModel{ + return &EmbeddingModel{ Config: cfg, Client: client, - } - return m, nil + }, nil } var _ embeddingmodels.EmbeddingModel = EmbeddingModel{} From cbcfa61673eb50401daf2988444fbf217ce32698 Mon Sep 17 00:00:00 2001 From: duwenxin Date: Mon, 2 Mar 2026 11:39:50 -0500 Subject: [PATCH 2/7] resolve comments and add tests --- .../configuration/embedding-models/gemini.md | 25 +++++++--------- internal/embeddingmodels/gemini/gemini.go | 11 +++---- .../embeddingmodels/gemini/gemini_test.go | 29 +++++++++++++++++-- 3 files changed, 41 insertions(+), 24 deletions(-) diff --git a/docs/en/documentation/configuration/embedding-models/gemini.md b/docs/en/documentation/configuration/embedding-models/gemini.md index b9d481b25a51..3be9eca03eab 100644 --- a/docs/en/documentation/configuration/embedding-models/gemini.md +++ b/docs/en/documentation/configuration/embedding-models/gemini.md @@ -51,7 +51,6 @@ more information. [modellist]: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings#supported-models - ## Example ### Using Google AI @@ -86,20 +85,18 @@ dimension: 768 [adc]: https://docs.cloud.google.com/docs/authentication/provide-credentials-adc - -{{< notice tip >}} - Use environment variable replacement with the format - ${ENV_NAME} instead of hardcoding your secrets into the configuration file. +{{< notice tip >}} Use environment variable replacement with the format +${ENV_NAME} instead of hardcoding your secrets into the configuration file. {{< /notice >}} ## Reference -| **field** | **type** | **required** | **description** | -| ----------- | :------: | :----------: | ------------------------------------------------------------- | -| type | string | true | Must be `gemini`. | -| model | string | true | The Gemini model ID to use (e.g., `text-embedding-004`). | -| useVertexAI | boolean | true | Set to `true` to use Vertex AI. Default is false (Google AI). | -| project | string | false | GCP Project ID (required if `useVertexAI` is `true`). | -| location | string | false | GCP Location (required if `useVertexAI` is `true`). | -| apiKey | string | false | Your API Key from Google AI Studio. | -| dimension | integer | false | The number of dimensions in the output vector (e.g., `768`). | +| **field** | **type** | **required** | **description** | +| ----------- | :------: | :----------: | ---------------------------------------------------------------------------------------------------------------------------------------------------- | +| type | string | true | Must be `gemini`. | +| model | string | true | The Gemini model ID to use (e.g., `text-embedding-004`). | +| useVertexAI | boolean | false | Set to `true` to use Vertex AI. Default is false (Google AI). | +| project | string | false | GCP Project ID (required if `useVertexAI` is `true`). | +| location | string | false | GCP Location (required if `useVertexAI` is `true`). | +| apiKey | string | false | Your API Key from Google AI Studio. Required if `useVertexAI` is `false` and not set via `GOOGLE_API_KEY` or `GEMINI_API_KEY` environment variables. | +| dimension | integer | false | The number of dimensions in the output vector (e.g., `768`). | diff --git a/internal/embeddingmodels/gemini/gemini.go b/internal/embeddingmodels/gemini/gemini.go index bd578b4946e3..50a99c02e091 100644 --- a/internal/embeddingmodels/gemini/gemini.go +++ b/internal/embeddingmodels/gemini/gemini.go @@ -34,7 +34,7 @@ type Config struct { Name string `yaml:"name" validate:"required"` Type string `yaml:"type" validate:"required"` Model string `yaml:"model" validate:"required"` - UseVertexAI bool `yaml:"useVertexAI" validate:"required"` + UseVertexAI bool `yaml:"useVertexAI"` ApiKey string `yaml:"apiKey"` Project string `yaml:"project"` Location string `yaml:"location"` @@ -54,9 +54,9 @@ func (cfg Config) Initialize(ctx context.Context) (embeddingmodels.EmbeddingMode apiKey := cfg.ApiKey if apiKey == "" { apiKey = os.Getenv("GOOGLE_API_KEY") - if apiKey == "" { - apiKey = os.Getenv("GEMINI_API_KEY") - } + } + if apiKey == "" { + apiKey = os.Getenv("GEMINI_API_KEY") } // Determine the Backend @@ -73,9 +73,6 @@ func (cfg Config) Initialize(ctx context.Context) (embeddingmodels.EmbeddingMode configs.Location = cfg.Location if configs.Location == "" { configs.Location = os.Getenv("GOOGLE_CLOUD_LOCATION") - if configs.Location == "" { - configs.Location = "us-central1" - } } if configs.Project == "" { diff --git a/internal/embeddingmodels/gemini/gemini_test.go b/internal/embeddingmodels/gemini/gemini_test.go index d3ca50a96d38..5bdca369f123 100644 --- a/internal/embeddingmodels/gemini/gemini_test.go +++ b/internal/embeddingmodels/gemini/gemini_test.go @@ -48,7 +48,7 @@ func TestParseFromYamlGemini(t *testing.T) { }, }, { - desc: "full example with optional fields", + desc: "full example with Google AI fields", in: ` kind: embeddingModel name: complex-gemini @@ -67,6 +67,30 @@ func TestParseFromYamlGemini(t *testing.T) { }, }, }, + { + desc: "Vertex AI configuration", + in: ` + kind: embeddingModels + name: vertex-gemini + type: gemini + model: text-embedding-004 + useVertexAI: true + project: "my-project" + location: "us-central1" + dimension: 512 + `, + want: map[string]embeddingmodels.EmbeddingModelConfig{ + "vertex-gemini": gemini.Config{ + Name: "vertex-gemini", + Type: gemini.EmbeddingModelType, + Model: "text-embedding-004", + UseVertexAI: true, + Project: "my-project", + Location: "us-central1", + Dimension: 512, + }, + }, + }, } for _, tc := range tcs { t.Run(tc.desc, func(t *testing.T) { @@ -81,6 +105,7 @@ func TestParseFromYamlGemini(t *testing.T) { }) } } + func TestFailParseFromYamlGemini(t *testing.T) { tcs := []struct { desc string @@ -94,7 +119,6 @@ func TestFailParseFromYamlGemini(t *testing.T) { name: bad-model type: gemini `, - // Removed the specific model name from the prefix to match your output err: "error unmarshaling embeddingModel: unable to parse as \"bad-model\": Key: 'Config.Model' Error:Field validation for 'Model' failed on the 'required' tag", }, { @@ -106,7 +130,6 @@ func TestFailParseFromYamlGemini(t *testing.T) { model: text-embedding-004 invalid_param: true `, - // Updated to match the specific line-starting format of your error output err: "error unmarshaling embeddingModel: unable to parse as \"bad-field\": [1:1] unknown field \"invalid_param\"\n> 1 | invalid_param: true\n ^\n 2 | model: text-embedding-004\n 3 | name: bad-field\n 4 | type: gemini", }, } From f5dfc52b0233f3ceaca736528e15238576197d2f Mon Sep 17 00:00:00 2001 From: duwenxin99 Date: Wed, 4 Mar 2026 14:56:16 -0500 Subject: [PATCH 3/7] update model version --- .../configuration/embedding-models/gemini.md | 4 ++-- internal/embeddingmodels/gemini/gemini_test.go | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/en/documentation/configuration/embedding-models/gemini.md b/docs/en/documentation/configuration/embedding-models/gemini.md index 3be9eca03eab..8bc5f3fe2297 100644 --- a/docs/en/documentation/configuration/embedding-models/gemini.md +++ b/docs/en/documentation/configuration/embedding-models/gemini.md @@ -76,7 +76,7 @@ how to set up ADC [here][adc]. kind: embeddingModels name: gemini-model type: gemini -model: text-embedding-004 +model: gemini-embedding-001 useVertexAI: true project: ${GOOGLE_CLOUD_PROJECT} location: us-central1 @@ -94,7 +94,7 @@ ${ENV_NAME} instead of hardcoding your secrets into the configuration file. | **field** | **type** | **required** | **description** | | ----------- | :------: | :----------: | ---------------------------------------------------------------------------------------------------------------------------------------------------- | | type | string | true | Must be `gemini`. | -| model | string | true | The Gemini model ID to use (e.g., `text-embedding-004`). | +| model | string | true | The Gemini model ID to use (e.g., `gemini-embedding-001`). | | useVertexAI | boolean | false | Set to `true` to use Vertex AI. Default is false (Google AI). | | project | string | false | GCP Project ID (required if `useVertexAI` is `true`). | | location | string | false | GCP Location (required if `useVertexAI` is `true`). | diff --git a/internal/embeddingmodels/gemini/gemini_test.go b/internal/embeddingmodels/gemini/gemini_test.go index 5bdca369f123..8b6a407bb959 100644 --- a/internal/embeddingmodels/gemini/gemini_test.go +++ b/internal/embeddingmodels/gemini/gemini_test.go @@ -43,7 +43,7 @@ func TestParseFromYamlGemini(t *testing.T) { "my-gemini-model": gemini.Config{ Name: "my-gemini-model", Type: gemini.EmbeddingModelType, - Model: "text-embedding-004", + Model: "gemini-embedding-001", }, }, }, @@ -53,7 +53,7 @@ func TestParseFromYamlGemini(t *testing.T) { kind: embeddingModel name: complex-gemini type: gemini - model: text-embedding-004 + model: gemini-embedding-001 apiKey: "test-api-key" dimension: 768 `, @@ -61,7 +61,7 @@ func TestParseFromYamlGemini(t *testing.T) { "complex-gemini": gemini.Config{ Name: "complex-gemini", Type: gemini.EmbeddingModelType, - Model: "text-embedding-004", + Model: "gemini-embedding-001", ApiKey: "test-api-key", Dimension: 768, }, @@ -73,7 +73,7 @@ func TestParseFromYamlGemini(t *testing.T) { kind: embeddingModels name: vertex-gemini type: gemini - model: text-embedding-004 + model: gemini-embedding-001 useVertexAI: true project: "my-project" location: "us-central1" @@ -83,7 +83,7 @@ func TestParseFromYamlGemini(t *testing.T) { "vertex-gemini": gemini.Config{ Name: "vertex-gemini", Type: gemini.EmbeddingModelType, - Model: "text-embedding-004", + Model: "gemini-embedding-001", UseVertexAI: true, Project: "my-project", Location: "us-central1", @@ -127,7 +127,7 @@ func TestFailParseFromYamlGemini(t *testing.T) { kind: embeddingModel name: bad-field type: gemini - model: text-embedding-004 + model: gemini-embedding-001 invalid_param: true `, err: "error unmarshaling embeddingModel: unable to parse as \"bad-field\": [1:1] unknown field \"invalid_param\"\n> 1 | invalid_param: true\n ^\n 2 | model: text-embedding-004\n 3 | name: bad-field\n 4 | type: gemini", From c8d597b8251d3b86bf6a7748621ad3303ed18f52 Mon Sep 17 00:00:00 2001 From: duwenxin99 Date: Thu, 26 Mar 2026 17:17:53 -0400 Subject: [PATCH 4/7] fix comments --- .../configuration/embedding-models/gemini.md | 2 +- internal/embeddingmodels/gemini/gemini.go | 38 +++++++++---------- .../embeddingmodels/gemini/gemini_test.go | 6 +-- 3 files changed, 21 insertions(+), 25 deletions(-) diff --git a/docs/en/documentation/configuration/embedding-models/gemini.md b/docs/en/documentation/configuration/embedding-models/gemini.md index 8bc5f3fe2297..7d5bad8779c1 100644 --- a/docs/en/documentation/configuration/embedding-models/gemini.md +++ b/docs/en/documentation/configuration/embedding-models/gemini.md @@ -73,7 +73,7 @@ Vertex AI uses Application Default Credentials (ADC) for authentication. Learn how to set up ADC [here][adc]. ```yaml -kind: embeddingModels +kind: embeddingModel name: gemini-model type: gemini model: gemini-embedding-001 diff --git a/internal/embeddingmodels/gemini/gemini.go b/internal/embeddingmodels/gemini/gemini.go index 50a99c02e091..40da0e8d0d89 100644 --- a/internal/embeddingmodels/gemini/gemini.go +++ b/internal/embeddingmodels/gemini/gemini.go @@ -59,29 +59,24 @@ func (cfg Config) Initialize(ctx context.Context) (embeddingmodels.EmbeddingMode apiKey = os.Getenv("GEMINI_API_KEY") } + // Try to resolve Project and Location + project := cfg.Project + if project == "" { + project = os.Getenv("GOOGLE_CLOUD_PROJECT") + } + + location := cfg.Location + if location == "" { + location = os.Getenv("GOOGLE_CLOUD_LOCATION") + } + // Determine the Backend - if cfg.UseVertexAI { + if project != "" && location != "" { // VertexAI API uses ADC for authentication. // ADC requires `Project` and `Location` to be set. configs.Backend = genai.BackendVertexAI - - configs.Project = cfg.Project - if configs.Project == "" { - configs.Project = os.Getenv("GOOGLE_CLOUD_PROJECT") - } - - configs.Location = cfg.Location - if configs.Location == "" { - configs.Location = os.Getenv("GOOGLE_CLOUD_LOCATION") - } - - if configs.Project == "" { - return nil, fmt.Errorf("vertex AI mode forced, but no project ID found in YAML or GOOGLE_CLOUD_PROJECT") - } - - if configs.Location == "" { - return nil, fmt.Errorf("vertex AI mode forced, but no location found in YAML or GOOGLE_CLOUD_LOCATION") - } + configs.Project = project + configs.Location = location } else if apiKey != "" { // Using Gemini API, which uses API Key for authentication. @@ -89,13 +84,14 @@ func (cfg Config) Initialize(ctx context.Context) (embeddingmodels.EmbeddingMode configs.APIKey = apiKey } else { - // Missing both credentials + // Missing credentials return nil, fmt.Errorf("missing credentials for Gemini embedding: " + "For Google AI: Provide 'apiKey' in YAML or set GOOGLE_API_KEY/GEMINI_API_KEY env vars. " + - "For Vertex AI: Set 'useVertexAI: true' and provide 'project'/'location' in YAML or via GOOGLE_CLOUD_PROJECT/GOOGLE_CLOUD_LOCATION env vars. " + + "For Vertex AI: Provide 'project'/'location' in YAML or via GOOGLE_CLOUD_PROJECT/GOOGLE_CLOUD_LOCATION env vars. " + "See documentation for details: https://googleapis.github.io/genai-toolbox/resources/embeddingmodels/gemini/") } + // Set user agent ua, err := util.UserAgentFromContext(ctx) if err != nil { return nil, fmt.Errorf("failed to get user agent from context: %w", err) diff --git a/internal/embeddingmodels/gemini/gemini_test.go b/internal/embeddingmodels/gemini/gemini_test.go index 8b6a407bb959..074e4bc9542a 100644 --- a/internal/embeddingmodels/gemini/gemini_test.go +++ b/internal/embeddingmodels/gemini/gemini_test.go @@ -37,7 +37,7 @@ func TestParseFromYamlGemini(t *testing.T) { kind: embeddingModel name: my-gemini-model type: gemini - model: text-embedding-004 + model: gemini-embedding-001 `, want: map[string]embeddingmodels.EmbeddingModelConfig{ "my-gemini-model": gemini.Config{ @@ -70,7 +70,7 @@ func TestParseFromYamlGemini(t *testing.T) { { desc: "Vertex AI configuration", in: ` - kind: embeddingModels + kind: embeddingModel name: vertex-gemini type: gemini model: gemini-embedding-001 @@ -130,7 +130,7 @@ func TestFailParseFromYamlGemini(t *testing.T) { model: gemini-embedding-001 invalid_param: true `, - err: "error unmarshaling embeddingModel: unable to parse as \"bad-field\": [1:1] unknown field \"invalid_param\"\n> 1 | invalid_param: true\n ^\n 2 | model: text-embedding-004\n 3 | name: bad-field\n 4 | type: gemini", + err: "error unmarshaling embeddingModel: unable to parse as \"bad-field\": [1:1] unknown field \"invalid_param\"\n> 1 | invalid_param: true\n ^\n 2 | model: gemini-embedding-001\n 3 | name: bad-field\n 4 | type: gemini", }, } for _, tc := range tcs { From 95f7d23551ff3637da04c51730e7f86bd193ff69 Mon Sep 17 00:00:00 2001 From: duwenxin99 Date: Thu, 26 Mar 2026 17:22:16 -0400 Subject: [PATCH 5/7] add logger --- internal/embeddingmodels/gemini/gemini.go | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/internal/embeddingmodels/gemini/gemini.go b/internal/embeddingmodels/gemini/gemini.go index 40da0e8d0d89..5a501c0355a2 100644 --- a/internal/embeddingmodels/gemini/gemini.go +++ b/internal/embeddingmodels/gemini/gemini.go @@ -50,6 +50,12 @@ func (cfg Config) EmbeddingModelConfigType() string { func (cfg Config) Initialize(ctx context.Context) (embeddingmodels.EmbeddingModel, error) { configs := &genai.ClientConfig{} + // Retrieve logger from context + l, err := util.LoggerFromContext(ctx) + if err != nil { + return nil, fmt.Errorf("unable to retrieve logger: %w", err) + } + // Get API Key apiKey := cfg.ApiKey if apiKey == "" { @@ -78,11 +84,15 @@ func (cfg Config) Initialize(ctx context.Context) (embeddingmodels.EmbeddingMode configs.Project = project configs.Location = location + l.InfoContext(ctx, "Using Vertex AI backend for Gemini embedding", "project", project, "location", location) + } else if apiKey != "" { // Using Gemini API, which uses API Key for authentication. configs.Backend = genai.BackendGeminiAPI configs.APIKey = apiKey + l.InfoContext(ctx, "Using Google AI (Gemini API) backend for Gemini embedding") + } else { // Missing credentials return nil, fmt.Errorf("missing credentials for Gemini embedding: " + From 1d334836ee8611c1e5795f586be3ce4bca6ad0a4 Mon Sep 17 00:00:00 2001 From: duwenxin99 Date: Thu, 26 Mar 2026 17:35:41 -0400 Subject: [PATCH 6/7] update tests --- .../configuration/embedding-models/gemini.md | 14 ++---- internal/embeddingmodels/gemini/gemini.go | 15 +++--- .../embeddingmodels/gemini/gemini_test.go | 50 ++++++++++++++----- 3 files changed, 48 insertions(+), 31 deletions(-) diff --git a/docs/en/documentation/configuration/embedding-models/gemini.md b/docs/en/documentation/configuration/embedding-models/gemini.md index 7d5bad8779c1..c0f8decedc79 100644 --- a/docs/en/documentation/configuration/embedding-models/gemini.md +++ b/docs/en/documentation/configuration/embedding-models/gemini.md @@ -16,13 +16,12 @@ high-dimensional vectors. Toolbox supports two authentication modes: -1. **Google AI (API Key):** Used if `useVertexAI` is `false` or unset. You must +1. **Google AI (API Key):** Used if you provide `apiKey` (or set `GOOGLE_API_KEY`/`GEMINI_API_KEY` environment variables). This uses the [Google AI Studio][ai-studio] backend. -2. **Vertex AI (ADC):** Used if `useVertexAI` is `true`. This uses [Application - Default Credentials (ADC)][adc]. When using this mode, you **must** specify - `project` and `location` (or set - `GOOGLE_CLOUD_PROJECT`/`GOOGLE_CLOUD_LOCATION` environment variables). +2. **Vertex AI (ADC):** Used if provided `project` and `location` (or set + `GOOGLE_CLOUD_PROJECT`/`GOOGLE_CLOUD_LOCATION` environment variables). This uses [Application + Default Credentials (ADC)][adc]. We recommend using an API key for quick testing and using Vertex AI with ADC for production environments. @@ -77,7 +76,6 @@ kind: embeddingModel name: gemini-model type: gemini model: gemini-embedding-001 -useVertexAI: true project: ${GOOGLE_CLOUD_PROJECT} location: us-central1 dimension: 768 @@ -95,8 +93,4 @@ ${ENV_NAME} instead of hardcoding your secrets into the configuration file. | ----------- | :------: | :----------: | ---------------------------------------------------------------------------------------------------------------------------------------------------- | | type | string | true | Must be `gemini`. | | model | string | true | The Gemini model ID to use (e.g., `gemini-embedding-001`). | -| useVertexAI | boolean | false | Set to `true` to use Vertex AI. Default is false (Google AI). | -| project | string | false | GCP Project ID (required if `useVertexAI` is `true`). | -| location | string | false | GCP Location (required if `useVertexAI` is `true`). | -| apiKey | string | false | Your API Key from Google AI Studio. Required if `useVertexAI` is `false` and not set via `GOOGLE_API_KEY` or `GEMINI_API_KEY` environment variables. | | dimension | integer | false | The number of dimensions in the output vector (e.g., `768`). | diff --git a/internal/embeddingmodels/gemini/gemini.go b/internal/embeddingmodels/gemini/gemini.go index 5a501c0355a2..f2e330d424b6 100644 --- a/internal/embeddingmodels/gemini/gemini.go +++ b/internal/embeddingmodels/gemini/gemini.go @@ -31,14 +31,13 @@ const EmbeddingModelType string = "gemini" var _ embeddingmodels.EmbeddingModelConfig = Config{} type Config struct { - Name string `yaml:"name" validate:"required"` - Type string `yaml:"type" validate:"required"` - Model string `yaml:"model" validate:"required"` - UseVertexAI bool `yaml:"useVertexAI"` - ApiKey string `yaml:"apiKey"` - Project string `yaml:"project"` - Location string `yaml:"location"` - Dimension int32 `yaml:"dimension"` + Name string `yaml:"name" validate:"required"` + Type string `yaml:"type" validate:"required"` + Model string `yaml:"model" validate:"required"` + ApiKey string `yaml:"apiKey"` + Project string `yaml:"project"` + Location string `yaml:"location"` + Dimension int32 `yaml:"dimension"` } // Returns the embedding model type diff --git a/internal/embeddingmodels/gemini/gemini_test.go b/internal/embeddingmodels/gemini/gemini_test.go index 074e4bc9542a..be58569f692e 100644 --- a/internal/embeddingmodels/gemini/gemini_test.go +++ b/internal/embeddingmodels/gemini/gemini_test.go @@ -16,6 +16,7 @@ package gemini_test import ( "context" + "strings" "testing" "github.com/google/go-cmp/cmp" @@ -74,20 +75,18 @@ func TestParseFromYamlGemini(t *testing.T) { name: vertex-gemini type: gemini model: gemini-embedding-001 - useVertexAI: true project: "my-project" location: "us-central1" dimension: 512 `, want: map[string]embeddingmodels.EmbeddingModelConfig{ "vertex-gemini": gemini.Config{ - Name: "vertex-gemini", - Type: gemini.EmbeddingModelType, - Model: "gemini-embedding-001", - UseVertexAI: true, - Project: "my-project", - Location: "us-central1", - Dimension: 512, + Name: "vertex-gemini", + Type: gemini.EmbeddingModelType, + Model: "gemini-embedding-001", + Project: "my-project", + Location: "us-central1", + Dimension: 512, }, }, }, @@ -132,15 +131,40 @@ func TestFailParseFromYamlGemini(t *testing.T) { `, err: "error unmarshaling embeddingModel: unable to parse as \"bad-field\": [1:1] unknown field \"invalid_param\"\n> 1 | invalid_param: true\n ^\n 2 | model: gemini-embedding-001\n 3 | name: bad-field\n 4 | type: gemini", }, + { + desc: "missing both Vertex and Google AI credentials", + in: ` + kind: embeddingModel + name: missing-creds + type: gemini + model: text-embedding-004 + `, + err: "unable to initialize embedding model \"missing-creds\": missing credentials for Gemini embedding: For Google AI: Provide 'apiKey' in YAML or set GOOGLE_API_KEY/GEMINI_API_KEY env vars. For Vertex AI: Provide 'project'/'location' in YAML or via GOOGLE_CLOUD_PROJECT/GOOGLE_CLOUD_LOCATION env vars. See documentation for details: https://googleapis.github.io/genai-toolbox/resources/embeddingmodels/gemini/", + }, } for _, tc := range tcs { t.Run(tc.desc, func(t *testing.T) { - _, _, _, _, _, _, err := server.UnmarshalResourceConfig(context.Background(), testutils.FormatYaml(tc.in)) - if err == nil { - t.Fatalf("expect parsing to fail") + t.Setenv("GOOGLE_API_KEY", "") + t.Setenv("GEMINI_API_KEY", "") + t.Setenv("GOOGLE_CLOUD_PROJECT", "") + t.Setenv("GOOGLE_CLOUD_LOCATION", "") + + _, embeddingConfigs, _, _, _, _, err := server.UnmarshalResourceConfig(context.Background(), testutils.FormatYaml(tc.in)) + if err != nil { + if err.Error() != tc.err { + t.Fatalf("unexpected unmarshal error:\ngot: %q\nwant: %q", err.Error(), tc.err) + } + return } - if err.Error() != tc.err { - t.Fatalf("unexpected error:\ngot: %q\nwant: %q", err.Error(), tc.err) + + for _, cfg := range embeddingConfigs { + _, err = cfg.Initialize() + if err == nil { + t.Fatalf("expect initialization to fail for case: %s", tc.desc) + } + if !strings.Contains(err.Error(), tc.err) { + t.Fatalf("unexpected init error:\ngot: %q\nwant: %q", err.Error(), tc.err) + } } }) } From 959a84066bb8df804170ae446200f2f2b32b0e0a Mon Sep 17 00:00:00 2001 From: duwenxin Date: Sun, 29 Mar 2026 23:21:27 -0400 Subject: [PATCH 7/7] fix lint --- internal/embeddingmodels/gemini/gemini_test.go | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/internal/embeddingmodels/gemini/gemini_test.go b/internal/embeddingmodels/gemini/gemini_test.go index be58569f692e..e6f8b1b5a3fd 100644 --- a/internal/embeddingmodels/gemini/gemini_test.go +++ b/internal/embeddingmodels/gemini/gemini_test.go @@ -151,19 +151,19 @@ func TestFailParseFromYamlGemini(t *testing.T) { _, embeddingConfigs, _, _, _, _, err := server.UnmarshalResourceConfig(context.Background(), testutils.FormatYaml(tc.in)) if err != nil { - if err.Error() != tc.err { - t.Fatalf("unexpected unmarshal error:\ngot: %q\nwant: %q", err.Error(), tc.err) - } - return + if err.Error() != tc.err { + t.Fatalf("unexpected unmarshal error:\ngot: %q\nwant: %q", err.Error(), tc.err) + } + return } for _, cfg := range embeddingConfigs { _, err = cfg.Initialize() if err == nil { - t.Fatalf("expect initialization to fail for case: %s", tc.desc) + t.Fatalf("expect initialization to fail for case: %s", tc.desc) } if !strings.Contains(err.Error(), tc.err) { - t.Fatalf("unexpected init error:\ngot: %q\nwant: %q", err.Error(), tc.err) + t.Fatalf("unexpected init error:\ngot: %q\nwant: %q", err.Error(), tc.err) } } })