Merge branch 'main' into fix-bigtable-resource-cleanup

Author: NirajNandre

docs/BIGQUERY_README.md

@@ -76,6 +76,7 @@ export BIGQUERY_PROJECT="<your-gcp-project-id>"
 export BIGQUERY_LOCATION="<your-dataset-location>"  # Optional
 export BIGQUERY_USE_CLIENT_OAUTH="true"  # Optional: true, false, or a custom header name
 export BIGQUERY_SCOPES="<comma-separated-scopes>"  # Optional
+export BIGQUERY_IMPERSONATE_SERVICE_ACCOUNT="<service-account-email>"  # Optional: Service account to impersonate
 ```
 
 Add the following configuration to your MCP client (e.g., `settings.json` for Gemini CLI, `mcp_config.json` for Antigravity):

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

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,17 @@ 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 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 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 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 +44,19 @@ 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 +66,31 @@ apiKey: ${GOOGLE_API_KEY}
 dimension: 768
 ```
 
-{{< notice tip >}}
-Use environment variable replacement with the format ${ENV_NAME}
-instead of hardcoding your secrets into the configuration file.
+### Using Vertex AI
+
+Vertex AI uses Application Default Credentials (ADC) for authentication. Learn
+how to set up ADC [here][adc].
+
+```yaml
+kind: embeddingModel
+name: gemini-model
+type: gemini
+model: gemini-embedding-001
+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.
 {{< /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., `gemini-embedding-001`).                                                                                             |
+| dimension   | integer  |    false     | The number of dimensions in the output vector (e.g., `768`).                                                                                         |

docs/en/integrations/bigquery/prebuilt-configs/bigquery.md

@@ -14,6 +14,10 @@ description: "Details of the BigQuery prebuilt configuration."
         OAuth access token for authentication. Defaults to `false`.
     *   `BIGQUERY_SCOPES`: (Optional) A comma-separated list of OAuth scopes to
         use for authentication.
+    *   `BIGQUERY_IMPERSONATE_SERVICE_ACCOUNT`: (Optional) Service account email
+        to impersonate when making BigQuery and Dataplex API calls. The
+        authenticated principal must have `roles/iam.serviceAccountTokenCreator`
+        on the target service account.
 *   **Permissions:**
     *   **BigQuery User** (`roles/bigquery.user`) to execute queries and view
         metadata.

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"
@@ -34,6 +35,8 @@ type Config struct {
 	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"`
 }
 
@@ -44,12 +47,60 @@ 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
+
+	// 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 == "" {
+		apiKey = os.Getenv("GOOGLE_API_KEY")
+	}
+	if apiKey == "" {
+		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 project != "" && location != "" {
+		// VertexAI API uses ADC for authentication.
+		// ADC requires `Project` and `Location` to be set.
+		configs.Backend = genai.BackendVertexAI
+		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: " +
+			"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/")
+	}
+
+	// Set user agent
 	ua, err := util.UserAgentFromContext(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get user agent from context: %w", err)
@@ -63,14 +114,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{}

internal/embeddingmodels/gemini/gemini_test.go

@@ -16,6 +16,7 @@ package gemini_test
 
 import (
 	"context"
+	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -37,36 +38,58 @@ 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{
 					Name:  "my-gemini-model",
 					Type:  gemini.EmbeddingModelType,
-					Model: "text-embedding-004",
+					Model: "gemini-embedding-001",
 				},
 			},
 		},
 		{
-			desc: "full example with optional fields",
+			desc: "full example with Google AI fields",
 			in: `
             kind: embeddingModel
             name: complex-gemini
             type: gemini
-            model: text-embedding-004
+            model: gemini-embedding-001
             apiKey: "test-api-key"
             dimension: 768
             `,
 			want: map[string]embeddingmodels.EmbeddingModelConfig{
 				"complex-gemini": gemini.Config{
 					Name:      "complex-gemini",
 					Type:      gemini.EmbeddingModelType,
-					Model:     "text-embedding-004",
+					Model:     "gemini-embedding-001",
 					ApiKey:    "test-api-key",
 					Dimension: 768,
 				},
 			},
 		},
+		{
+			desc: "Vertex AI configuration",
+			in: `
+            kind: embeddingModel
+            name: vertex-gemini
+            type: gemini
+            model: gemini-embedding-001
+            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",
+					Project:   "my-project",
+					Location:  "us-central1",
+					Dimension: 512,
+				},
+			},
+		},
 	}
 	for _, tc := range tcs {
 		t.Run(tc.desc, func(t *testing.T) {
@@ -81,6 +104,7 @@ func TestParseFromYamlGemini(t *testing.T) {
 		})
 	}
 }
+
 func TestFailParseFromYamlGemini(t *testing.T) {
 	tcs := []struct {
 		desc string
@@ -94,7 +118,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",
 		},
 		{
@@ -103,21 +126,45 @@ func TestFailParseFromYamlGemini(t *testing.T) {
             kind: embeddingModel
             name: bad-field
             type: gemini
-            model: text-embedding-004
+            model: gemini-embedding-001
             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",
+			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)
+				}
 			}
 		})
 	}

internal/prebuiltconfigs/tools/bigquery.yaml

@@ -20,6 +20,7 @@ sources:
         useClientOAuth: ${BIGQUERY_USE_CLIENT_OAUTH:false}
         scopes: ${BIGQUERY_SCOPES:}
         maxQueryResultRows: ${BIGQUERY_MAX_QUERY_RESULT_ROWS:50}
+        impersonateServiceAccount: ${BIGQUERY_IMPERSONATE_SERVICE_ACCOUNT:}
 tools:
     analyze_contribution:
         kind: bigquery-analyze-contribution