feat: support native Gemini text embeddings

Implements CreateEmbedding/CreateBatchEmbedding via batchEmbedContents
(Vertex single-input path), output_dimensionality provider_opt, and
ordering/cardinality/dimension validation. Adds docs, schema update,
and examples/rag/gemini_embeddings.yaml.

Assisted-By: docker-agent
Signed-off-by: David Gageot <david.gageot@docker.com>

agent-schema.json

@@ -1639,7 +1639,7 @@
         },
         "provider_opts": {
           "type": "object",
-          "description": "Provider-specific options. Sampling parameters: top_k (integer, supported by anthropic, google, amazon-bedrock, and custom OpenAI-compatible providers like vLLM/Ollama), repetition_penalty (float, forwarded to custom OpenAI-compatible providers), min_p (float, forwarded to custom providers), seed (integer, forwarded to OpenAI). Lifecycle: unload_api (string) overrides the unload endpoint inherited from the provider config (relative path resolved against base_url's scheme+host, or an absolute URL); used by the runtime's `unload` on_agent_switch builtin hook to release model resources between agent switches. Infrastructure options: http_headers (map of string to string, adds custom HTTP headers to every request; used for OpenAI-compatible providers like github-copilot which requires Copilot-Integration-Id). dmr: runtime_flags; supports_images (boolean) and supports_pdf (boolean) declare which document attachment types the local model accepts. DMR-hosted models are not in the models.dev catalog, so attachment capabilities cannot be detected automatically and default to text-only; set supports_images: true for a vision model so image attachments are forwarded instead of silently dropped. anthropic/amazon-bedrock (Claude): interleaved_thinking (boolean, default true). anthropic: thinking_display ('summarized', 'omitted', or 'display') controls whether thinking blocks are returned in responses when thinking is enabled. Claude Opus 4.7 hides thinking by default ('omitted'); set thinking_display: summarized to receive thinking blocks. 'display' (full thinking blocks) is only accepted by pre-4.6 token-thinking models (e.g. Sonnet 4.5, Haiku 4.5); models from the adaptive-thinking generation onward (Opus/Sonnet 4.6+, Sonnet 5, Fable 5) reject it, and docker-agent fails fast with a configuration error (fallbacks models are validated too); fallbacks (array of model ID strings, in priority order; enables the server-side-fallback beta so requests refused by safety classifiers, e.g. on Claude Fable 5, are retried with the listed models in a single round trip). OpenAI-compatible APIs: service_tier (string, forwarded unchanged to Chat Completions and Responses, including WebSocket and reranking; 'fast' or 'priority' enables Fast mode at premium pricing on supported models; omitted or empty leaves the API default, non-string values are ignored). openai/chatgpt: preserve_reasoning (boolean, default false) preserves ordered Responses output and encrypted reasoning for stateless continuation; native_tool_search also enables output replay. openai: cache_diagnostics (boolean, default false) logs Responses prompt-cache comparison results with --debug on GPT-5.6+; comparisons are scoped to the current session and resolved model; ignored for chatgpt and custom base_url endpoints. transport ('sse' or 'websocket') to choose between SSE and WebSocket streaming for the Responses API; native_tool_search (boolean, default false; also enables ordered output replay) declares the agent's deferred tools with defer_loading and a server-executed tool_search tool so OpenAI's hosted tool search discovers them (only openai models supporting deferred tools, gpt-5.4+; ignored on chatgpt, OpenAI-compatible endpoints including a custom base_url, and Chat Completions requests). openai/anthropic/google: rerank_prompt (string) to fully override the system prompt used for RAG reranking (advanced - prefer using results.reranking.criteria for domain-specific guidance). Google: google_search (boolean) enables Google Search grounding, google_maps (boolean) enables Google Maps grounding, code_execution (boolean) enables server-side code execution.",
+          "description": "Provider-specific options. Sampling parameters: top_k (integer, supported by anthropic, google, amazon-bedrock, and custom OpenAI-compatible providers like vLLM/Ollama), repetition_penalty (float, forwarded to custom OpenAI-compatible providers), min_p (float, forwarded to custom providers), seed (integer, forwarded to OpenAI). Lifecycle: unload_api (string) overrides the unload endpoint inherited from the provider config (relative path resolved against base_url's scheme+host, or an absolute URL); used by the runtime's `unload` on_agent_switch builtin hook to release model resources between agent switches. Infrastructure options: http_headers (map of string to string, adds custom HTTP headers to every request; used for OpenAI-compatible providers like github-copilot which requires Copilot-Integration-Id). dmr: runtime_flags; supports_images (boolean) and supports_pdf (boolean) declare which document attachment types the local model accepts. DMR-hosted models are not in the models.dev catalog, so attachment capabilities cannot be detected automatically and default to text-only; set supports_images: true for a vision model so image attachments are forwarded instead of silently dropped. anthropic/amazon-bedrock (Claude): interleaved_thinking (boolean, default true). anthropic: thinking_display ('summarized', 'omitted', or 'display') controls whether thinking blocks are returned in responses when thinking is enabled. Claude Opus 4.7 hides thinking by default ('omitted'); set thinking_display: summarized to receive thinking blocks. 'display' (full thinking blocks) is only accepted by pre-4.6 token-thinking models (e.g. Sonnet 4.5, Haiku 4.5); models from the adaptive-thinking generation onward (Opus/Sonnet 4.6+, Sonnet 5, Fable 5) reject it, and docker-agent fails fast with a configuration error (fallbacks models are validated too); fallbacks (array of model ID strings, in priority order; enables the server-side-fallback beta so requests refused by safety classifiers, e.g. on Claude Fable 5, are retried with the listed models in a single round trip). OpenAI-compatible APIs: service_tier (string, forwarded unchanged to Chat Completions and Responses, including WebSocket and reranking; 'fast' or 'priority' enables Fast mode at premium pricing on supported models; omitted or empty leaves the API default, non-string values are ignored). openai/chatgpt: preserve_reasoning (boolean, default false) preserves ordered Responses output and encrypted reasoning for stateless continuation; native_tool_search also enables output replay. openai: cache_diagnostics (boolean, default false) logs Responses prompt-cache comparison results with --debug on GPT-5.6+; comparisons are scoped to the current session and resolved model; ignored for chatgpt and custom base_url endpoints. transport ('sse' or 'websocket') to choose between SSE and WebSocket streaming for the Responses API; native_tool_search (boolean, default false; also enables ordered output replay) declares the agent's deferred tools with defer_loading and a server-executed tool_search tool so OpenAI's hosted tool search discovers them (only openai models supporting deferred tools, gpt-5.4+; ignored on chatgpt, OpenAI-compatible endpoints including a custom base_url, and Chat Completions requests). openai/anthropic/google: rerank_prompt (string) to fully override the system prompt used for RAG reranking (advanced - prefer using results.reranking.criteria for domain-specific guidance). Google: google_search (boolean) enables Google Search grounding, google_maps (boolean) enables Google Maps grounding, code_execution (boolean) enables server-side code execution, output_dimensionality (positive integer) truncates Gemini embedding vectors to that size when the model is used as a RAG embedding_model.",
           "additionalProperties": true
         },
         "track_usage": {

docs/providers/google/index.md

@@ -180,6 +180,47 @@ models:
 | `google_maps`    | Enables Google Maps grounding for location queries   |
 | `code_execution` | Enables server-side code execution for computations  |
 
+## Embeddings
+
+Gemini embedding models can back the [RAG](../../tools/rag/index.md)
+`chunked-embeddings` and `semantic-embeddings` strategies. Text is embedded
+natively through the Gemini API (`batchEmbedContents`) or Vertex AI, and vectors
+come back in input order. On Vertex AI, Gemini Embedding 2 models accept one
+input per request, so batches are sent one text at a time there.
+
+```yaml
+models:
+  gemini-embed:
+    provider: google
+    model: gemini-embedding-2
+    provider_opts:
+      output_dimensionality: 768   # optional; defaults to the model's native size
+
+rag:
+  docs:
+    docs: [./docs]
+    strategies:
+      - type: chunked-embeddings
+        embedding_model: gemini-embed
+        database: ./docs.db
+        vector_dimensions: 768       # must match output_dimensionality (or the native size)
+```
+
+| Option                  | Description                                                                                                           |
+| ----------------------- | --------------------------------------------------------------------------------------------------------------------- |
+| `output_dimensionality` | Positive integer. Truncates vectors to this size on models that support it (Matryoshka). Omit to use the native size. |
+
+`gemini-embedding-001` works the same way. Vectors from different models (or
+different `output_dimensionality` values) live in different vector spaces, so
+rebuild the index after changing either.
+
+Only text is embedded, and no `task_type` is sent: queries and documents share
+one vector space with no automatic task prefixes, and the same configuration
+works with Gemini Embedding 2, which rejects `task_type`. The Gemini Developer
+API does not report token usage for embeddings, so RAG cost statistics stay at
+zero; Vertex AI per-input token statistics are summed into the input token
+count when present.
+
 ## Vertex AI Model Garden
 
 You can use non-Gemini models (e.g. Claude, Llama) hosted on Google Cloud's

docs/tools/rag/index.md

@@ -67,6 +67,8 @@ strategies:
       overlap: 100
 ```
 
+Supported embedding providers: **OpenAI** (and OpenAI-compatible custom providers), **DMR**, **Google Gemini** (see [Embeddings](../../providers/google/index.md#embeddings)). `vector_dimensions` must match the model's output size, e.g. 1536 for `openai/text-embedding-3-small` or 3072 for `google/gemini-embedding-2`.
+
 ### Semantic Embeddings (LLM-Enhanced)
 
 Uses an LLM to generate semantic summaries of each chunk before embedding, capturing meaning and intent. Best for code search and understanding implementations.

examples/README.md

@@ -112,6 +112,7 @@ Examples that wire up one of the toolsets shipped with docker-agent
 | [`rag/hybrid.yaml`](rag/hybrid.yaml) | Hybrid BM25 + embeddings retrieval. |
 | [`rag/reranking.yaml`](rag/reranking.yaml) | Hybrid retrieval with a re-ranking model on top. |
 | [`rag/custom_provider.yaml`](rag/custom_provider.yaml) | Pointing the embeddings model at a custom OpenAI-compatible provider. |
+| [`rag/gemini_embeddings.yaml`](rag/gemini_embeddings.yaml) | Google Gemini embedding model with `output_dimensionality`. |
 
 ---
 

examples/rag/gemini_embeddings.yaml

@@ -0,0 +1,46 @@
+# This example demonstrates using a Google Gemini embedding model for RAG.
+#
+# Requires GOOGLE_API_KEY (or GEMINI_API_KEY). Vertex AI works too via
+# GOOGLE_GENAI_USE_VERTEXAI / project + location provider_opts.
+#
+# Only text is embedded, and no task_type is sent: queries and documents share
+# the same vector space without automatic task prefixes.
+
+models:
+  gemini-embed:
+    provider: google
+    # gemini-embedding-001 also works; vectors from different models are not
+    # comparable, so rebuild the index when switching.
+    model: gemini-embedding-2
+    provider_opts:
+      # Optional: truncate vectors (Matryoshka). Omit to keep the native size (3072).
+      # Must match vector_dimensions below.
+      output_dimensionality: 768
+
+agents:
+  root:
+    model: google/gemini-2.5-flash
+    description: assistant with RAG backed by Gemini embeddings
+    instruction: |
+      You are a helpful assistant with access to a knowledge base.
+      Use the search tool to find relevant information before answering.
+    toolsets:
+      - type: rag
+        ref: knowledge_base
+
+rag:
+  knowledge_base:
+    tool:
+      description: search the knowledge base for relevant information
+    docs:
+      - ./docs
+    strategies:
+      - type: chunked-embeddings
+        embedding_model: gemini-embed
+        database: ./gemini_embeddings.db
+        vector_dimensions: 768
+        chunking:
+          size: 1000
+          overlap: 100
+    results:
+      limit: 5

pkg/model/provider/gemini/embed.go

@@ -0,0 +1,153 @@
+package gemini
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"log/slog"
+	"math"
+	"strings"
+
+	"google.golang.org/genai"
+
+	"github.com/docker/docker-agent/pkg/model/provider/base"
+	"github.com/docker/docker-agent/pkg/model/provider/providerutil"
+)
+
+// outputDimensionalityOpt is the provider_opts key mapped to Gemini's outputDimensionality.
+const outputDimensionalityOpt = "output_dimensionality"
+
+// CreateEmbedding generates an embedding vector for the given text.
+func (c *Client) CreateEmbedding(ctx context.Context, text string) (*base.EmbeddingResult, error) {
+	batch, err := c.CreateBatchEmbedding(ctx, []string{text})
+	if err != nil {
+		return nil, err
+	}
+	if len(batch.Embeddings) == 0 {
+		return nil, errors.New("no embedding returned from Gemini")
+	}
+	return &base.EmbeddingResult{
+		Embedding:   batch.Embeddings[0],
+		InputTokens: batch.InputTokens,
+		TotalTokens: batch.TotalTokens,
+		Cost:        batch.Cost,
+	}, nil
+}
+
+// CreateBatchEmbedding generates embedding vectors for multiple texts, in input order.
+func (c *Client) CreateBatchEmbedding(ctx context.Context, texts []string) (*base.BatchEmbeddingResult, error) {
+	if len(texts) == 0 {
+		return &base.BatchEmbeddingResult{Embeddings: [][]float64{}}, nil
+	}
+
+	config, err := c.embedConfig()
+	if err != nil {
+		return nil, err
+	}
+
+	client, err := c.clientFn(ctx)
+	if err != nil {
+		slog.ErrorContext(ctx, "Failed to create Gemini client for embedding", "error", err)
+		return nil, err
+	}
+
+	slog.DebugContext(ctx, "Creating Gemini embeddings", "model", c.ModelConfig.Model, "batch_size", len(texts))
+
+	contents := make([]*genai.Content, len(texts))
+	for i, text := range texts {
+		contents[i] = genai.NewContentFromText(text, genai.RoleUser)
+	}
+
+	var embeddings []*genai.ContentEmbedding
+	if c.vertexEmbedsOneContentPerRequest() {
+		embeddings, err = c.embedEach(ctx, client, contents, config)
+	} else {
+		embeddings, err = c.embed(ctx, client, contents, config)
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	if len(embeddings) != len(texts) {
+		return nil, fmt.Errorf("expected %d embeddings, got %d", len(texts), len(embeddings))
+	}
+
+	result := &base.BatchEmbeddingResult{Embeddings: make([][]float64, len(embeddings))}
+	for i, e := range embeddings {
+		if e == nil || len(e.Values) == 0 {
+			return nil, fmt.Errorf("gemini returned an empty embedding at index %d", i)
+		}
+		if dims := len(embeddings[0].Values); len(e.Values) != dims {
+			return nil, fmt.Errorf("gemini returned inconsistent embedding dimensions: index 0 has %d, index %d has %d", dims, i, len(e.Values))
+		}
+		vec := make([]float64, len(e.Values))
+		for j, v := range e.Values {
+			vec[j] = float64(v)
+		}
+		result.Embeddings[i] = vec
+		if e.Statistics != nil {
+			result.InputTokens += int64(math.Round(float64(e.Statistics.TokenCount)))
+		}
+	}
+	// Embeddings have no output tokens; Gemini Developer API reports no usage at all.
+	result.TotalTokens = result.InputTokens
+
+	slog.DebugContext(ctx, "Gemini embeddings created",
+		"batch_size", len(result.Embeddings),
+		"dimension", len(result.Embeddings[0]),
+		"input_tokens", result.InputTokens)
+
+	return result, nil
+}
+
+// embed sends all contents in a single request.
+func (c *Client) embed(ctx context.Context, client *genai.Client, contents []*genai.Content, config *genai.EmbedContentConfig) ([]*genai.ContentEmbedding, error) {
+	resp, err := client.Models.EmbedContent(ctx, c.ModelConfig.Model, contents, config)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create embeddings: %w", wrapGeminiError(err))
+	}
+	if resp == nil {
+		return nil, errors.New("gemini embedding response is nil")
+	}
+	return resp.Embeddings, nil
+}
+
+// embedEach sends one request per content. Each response must carry exactly
+// one vector, otherwise texts and vectors would silently misalign.
+func (c *Client) embedEach(ctx context.Context, client *genai.Client, contents []*genai.Content, config *genai.EmbedContentConfig) ([]*genai.ContentEmbedding, error) {
+	embeddings := make([]*genai.ContentEmbedding, len(contents))
+	for i, content := range contents {
+		batch, err := c.embed(ctx, client, []*genai.Content{content}, config)
+		if err != nil {
+			return nil, err
+		}
+		if len(batch) != 1 {
+			return nil, fmt.Errorf("expected 1 embedding for input %d, got %d", i, len(batch))
+		}
+		embeddings[i] = batch[0]
+	}
+	return embeddings, nil
+}
+
+// vertexEmbedsOneContentPerRequest mirrors the genai SDK routing: on Vertex AI,
+// Gemini embedding models other than gemini-embedding-001 use embedContent,
+// which accepts a single content per request.
+func (c *Client) vertexEmbedsOneContentPerRequest() bool {
+	model := c.ModelConfig.Model
+	return c.apiSurface == apiSurfaceVertexAI && strings.Contains(model, "gemini") && model != "gemini-embedding-001"
+}
+
+// embedConfig builds the embedding request config from provider_opts.
+// task_type is deliberately not sent: Gemini Embedding 2 rejects it.
+func (c *Client) embedConfig() (*genai.EmbedContentConfig, error) {
+	config := &genai.EmbedContentConfig{}
+	if _, set := c.ModelConfig.ProviderOpts[outputDimensionalityOpt]; !set {
+		return config, nil
+	}
+	dims, ok := providerutil.GetProviderOptInt64(c.ModelConfig.ProviderOpts, outputDimensionalityOpt)
+	if !ok || dims <= 0 || dims > math.MaxInt32 {
+		return nil, fmt.Errorf("provider_opts.%s must be a positive integer, got %v", outputDimensionalityOpt, c.ModelConfig.ProviderOpts[outputDimensionalityOpt])
+	}
+	config.OutputDimensionality = new(int32(dims))
+	return config, nil
+}

pkg/model/provider/gemini/embed_test.go

@@ -0,0 +1,441 @@
+package gemini
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/genai"
+
+	"github.com/docker/docker-agent/pkg/config/latest"
+	"github.com/docker/docker-agent/pkg/environment"
+	"github.com/docker/docker-agent/pkg/model/provider/base"
+	"github.com/docker/docker-agent/pkg/model/provider/options"
+	"github.com/docker/docker-agent/pkg/modelerrors"
+)
+
+// batchEmbedRequest mirrors the Gemini Developer API batchEmbedContents body.
+type batchEmbedRequest struct {
+	Requests []struct {
+		Model   string `json:"model"`
+		Content struct {
+			Parts []struct {
+				Text string `json:"text"`
+			} `json:"parts"`
+		} `json:"content"`
+		TaskType             *string `json:"taskType"`
+		OutputDimensionality *int    `json:"outputDimensionality"`
+	} `json:"requests"`
+}
+
+func newEmbedClient(t *testing.T, baseURL string, providerOpts map[string]any, opts ...options.Opt) *Client {
+	t.Helper()
+	return newEmbedClientForModel(t, "gemini-embedding-001", baseURL, providerOpts, opts...)
+}
+
+func newEmbedClientForModel(t *testing.T, model, baseURL string, providerOpts map[string]any, opts ...options.Opt) *Client {
+	t.Helper()
+	cfg := &latest.ModelConfig{
+		Provider:     "google",
+		Model:        model,
+		BaseURL:      baseURL,
+		ProviderOpts: providerOpts,
+	}
+	env := environment.NewMapEnvProvider(map[string]string{
+		"GOOGLE_API_KEY":                  "test-key",
+		environment.DockerDesktopTokenEnv: "test-dd-token",
+	})
+	client, err := NewClient(t.Context(), cfg, env, opts...)
+	require.NoError(t, err)
+	return client
+}
+
+// newVertexEmbedClient routes through Vertex AI. With a custom BaseURL and no
+// project, location or API key, the genai SDK skips ADC lookup.
+func newVertexEmbedClient(t *testing.T, model, baseURL string, providerOpts map[string]any) *Client {
+	t.Helper()
+	cfg := &latest.ModelConfig{
+		Provider:     "google",
+		Model:        model,
+		BaseURL:      baseURL,
+		ProviderOpts: providerOpts,
+	}
+	env := environment.NewMapEnvProvider(map[string]string{"GOOGLE_GENAI_USE_VERTEXAI": "1"})
+	client, err := NewClient(t.Context(), cfg, env)
+	require.NoError(t, err)
+	require.Equal(t, apiSurfaceVertexAI, client.apiSurface)
+	return client
+}
+
+// recordingServer captures every request path and JSON body, replying with
+// responses[i] to the i-th request.
+func recordingServer(t *testing.T, responses ...string) (*httptest.Server, *[]string, *[]map[string]any) {
+	t.Helper()
+	var (
+		mu     sync.Mutex
+		paths  []string
+		bodies []map[string]any
+	)
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var body map[string]any
+		assert.NoError(t, json.NewDecoder(r.Body).Decode(&body))
+		mu.Lock()
+		paths = append(paths, r.URL.Path)
+		bodies = append(bodies, body)
+		n := len(paths)
+		mu.Unlock()
+		if !assert.LessOrEqual(t, n, len(responses), "unexpected extra request") {
+			w.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = io.WriteString(w, responses[n-1])
+	}))
+	t.Cleanup(srv.Close)
+	return srv, &paths, &bodies
+}
+
+func embedServer(t *testing.T, status int, body string, captured *batchEmbedRequest, calls *atomic.Int32) *httptest.Server {
+	t.Helper()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if calls != nil {
+			calls.Add(1)
+		}
+		assert.Equal(t, http.MethodPost, r.Method)
+		assert.Equal(t, "/v1beta/models/gemini-embedding-001:batchEmbedContents", r.URL.Path)
+		if captured != nil {
+			raw, err := io.ReadAll(r.Body)
+			assert.NoError(t, err)
+			assert.NoError(t, json.Unmarshal(raw, captured))
+		}
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(status)
+		_, _ = io.WriteString(w, body)
+	}))
+	t.Cleanup(srv.Close)
+	return srv
+}
+
+func TestCreateBatchEmbedding_WireShapeAndOrder(t *testing.T) {
+	t.Parallel()
+
+	for _, gateway := range []bool{false, true} {
+		t.Run(fmt.Sprintf("gateway=%t", gateway), func(t *testing.T) {
+			t.Parallel()
+
+			var captured batchEmbedRequest
+			srv := embedServer(t, http.StatusOK, `{"embeddings":[{"values":[0.1,0.2,0.3]},{"values":[0.4,0.5,0.6]}]}`, &captured, nil)
+			var client *Client
+			if gateway {
+				client = newEmbedClient(t, "", nil, options.WithGateway(srv.URL))
+			} else {
+				client = newEmbedClient(t, srv.URL, nil)
+			}
+
+			result, err := client.CreateBatchEmbedding(t.Context(), []string{"first", "second"})
+			require.NoError(t, err)
+
+			require.Len(t, captured.Requests, 2)
+			assert.Equal(t, "models/gemini-embedding-001", captured.Requests[0].Model)
+			assert.Equal(t, "first", captured.Requests[0].Content.Parts[0].Text)
+			assert.Equal(t, "second", captured.Requests[1].Content.Parts[0].Text)
+			for _, req := range captured.Requests {
+				assert.Nil(t, req.TaskType, "task_type must not be sent by default")
+				assert.Nil(t, req.OutputDimensionality)
+			}
+
+			require.Len(t, result.Embeddings, 2)
+			assert.InDeltaSlice(t, []float64{0.1, 0.2, 0.3}, result.Embeddings[0], 1e-6)
+			assert.InDeltaSlice(t, []float64{0.4, 0.5, 0.6}, result.Embeddings[1], 1e-6)
+			assert.Zero(t, result.InputTokens, "Gemini Developer API reports no embedding usage")
+			assert.Zero(t, result.TotalTokens)
+			assert.Zero(t, result.Cost)
+		})
+	}
+}
+
+func TestCreateBatchEmbedding_OutputDimensionality(t *testing.T) {
+	t.Parallel()
+
+	var captured batchEmbedRequest
+	srv := embedServer(t, http.StatusOK, `{"embeddings":[{"values":[0.1,0.2]}]}`, &captured, nil)
+	client := newEmbedClient(t, srv.URL, map[string]any{"output_dimensionality": 768})
+
+	_, err := client.CreateBatchEmbedding(t.Context(), []string{"text"})
+	require.NoError(t, err)
+
+	require.Len(t, captured.Requests, 1)
+	require.NotNil(t, captured.Requests[0].OutputDimensionality)
+	assert.Equal(t, 768, *captured.Requests[0].OutputDimensionality)
+}
+
+func TestCreateBatchEmbedding_InvalidOutputDimensionality(t *testing.T) {
+	t.Parallel()
+
+	for name, value := range map[string]any{"zero": 0, "negative": -1, "string": "768", "float": 1.5} {
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			var calls atomic.Int32
+			srv := embedServer(t, http.StatusOK, `{}`, nil, &calls)
+			client := newEmbedClient(t, srv.URL, map[string]any{"output_dimensionality": value})
+
+			_, err := client.CreateBatchEmbedding(t.Context(), []string{"text"})
+			require.ErrorContains(t, err, "output_dimensionality")
+			assert.Zero(t, calls.Load(), "invalid config must fail before any request")
+		})
+	}
+}
+
+func TestCreateBatchEmbedding_UsageFromStatistics(t *testing.T) {
+	t.Parallel()
+
+	srv := embedServer(t, http.StatusOK, `{"embeddings":[{"values":[0.1],"statistics":{"tokenCount":3}},{"values":[0.2],"statistics":{"tokenCount":4}}],"metadata":{"billableCharacterCount":42}}`, nil, nil)
+	client := newEmbedClient(t, srv.URL, nil)
+
+	result, err := client.CreateBatchEmbedding(t.Context(), []string{"a", "b"})
+	require.NoError(t, err)
+	assert.Equal(t, int64(7), result.InputTokens)
+	assert.Equal(t, int64(7), result.TotalTokens)
+}
+
+func TestCreateEmbedding_ReturnsFirstVector(t *testing.T) {
+	t.Parallel()
+
+	var captured batchEmbedRequest
+	srv := embedServer(t, http.StatusOK, `{"embeddings":[{"values":[1,2]}]}`, &captured, nil)
+	client := newEmbedClient(t, srv.URL, nil)
+
+	result, err := client.CreateEmbedding(t.Context(), "hello")
+	require.NoError(t, err)
+	require.Len(t, captured.Requests, 1)
+	assert.Equal(t, "hello", captured.Requests[0].Content.Parts[0].Text)
+	assert.Equal(t, []float64{1, 2}, result.Embedding)
+}
+
+func TestCreateBatchEmbedding_EmptyInputSkipsRequest(t *testing.T) {
+	t.Parallel()
+
+	var calls atomic.Int32
+	srv := embedServer(t, http.StatusOK, `{}`, nil, &calls)
+	client := newEmbedClient(t, srv.URL, nil)
+
+	result, err := client.CreateBatchEmbedding(t.Context(), nil)
+	require.NoError(t, err)
+	assert.Equal(t, &base.BatchEmbeddingResult{Embeddings: [][]float64{}}, result)
+	assert.Zero(t, calls.Load())
+}
+
+func TestCreateBatchEmbedding_CardinalityMismatch(t *testing.T) {
+	t.Parallel()
+
+	srv := embedServer(t, http.StatusOK, `{"embeddings":[{"values":[0.1]}]}`, nil, nil)
+	client := newEmbedClient(t, srv.URL, nil)
+
+	_, err := client.CreateBatchEmbedding(t.Context(), []string{"a", "b"})
+	require.ErrorContains(t, err, "expected 2 embeddings, got 1")
+}
+
+func TestCreateBatchEmbedding_EmptyVector(t *testing.T) {
+	t.Parallel()
+
+	srv := embedServer(t, http.StatusOK, `{"embeddings":[{"values":[0.1]},{}]}`, nil, nil)
+	client := newEmbedClient(t, srv.URL, nil)
+
+	_, err := client.CreateBatchEmbedding(t.Context(), []string{"a", "b"})
+	require.ErrorContains(t, err, "empty embedding at index 1")
+}
+
+func TestCreateBatchEmbedding_InconsistentDimensions(t *testing.T) {
+	t.Parallel()
+
+	srv := embedServer(t, http.StatusOK, `{"embeddings":[{"values":[0.1,0.2]},{"values":[0.3,0.4]},{"values":[0.5]}]}`, nil, nil)
+	client := newEmbedClient(t, srv.URL, nil)
+
+	_, err := client.CreateBatchEmbedding(t.Context(), []string{"a", "b", "c"})
+	require.ErrorContains(t, err, "inconsistent embedding dimensions: index 0 has 2, index 2 has 1")
+}
+
+func TestCreateBatchEmbedding_Embedding2DeveloperAPI(t *testing.T) {
+	t.Parallel()
+
+	srv, paths, bodies := recordingServer(t, `{"embeddings":[{"values":[0.1,0.2]},{"values":[0.3,0.4]}]}`)
+	client := newEmbedClientForModel(t, "gemini-embedding-2", srv.URL, map[string]any{"output_dimensionality": 2})
+
+	result, err := client.CreateBatchEmbedding(t.Context(), []string{"first", "second"})
+	require.NoError(t, err)
+
+	require.Equal(t, []string{"/v1beta/models/gemini-embedding-2:batchEmbedContents"}, *paths)
+	requests, ok := (*bodies)[0]["requests"].([]any)
+	require.True(t, ok)
+	require.Len(t, requests, 2)
+	for i, text := range []string{"first", "second"} {
+		req := requests[i].(map[string]any)
+		assert.Equal(t, "models/gemini-embedding-2", req["model"])
+		assert.Equal(t, text, req["content"].(map[string]any)["parts"].([]any)[0].(map[string]any)["text"])
+		assert.InDelta(t, 2, req["outputDimensionality"], 0)
+		assert.NotContains(t, req, "taskType")
+	}
+
+	require.Len(t, result.Embeddings, 2)
+	assert.InDeltaSlice(t, []float64{0.1, 0.2}, result.Embeddings[0], 1e-6)
+	assert.InDeltaSlice(t, []float64{0.3, 0.4}, result.Embeddings[1], 1e-6)
+	assert.Zero(t, result.InputTokens)
+}
+
+func TestCreateBatchEmbedding_VertexPredict(t *testing.T) {
+	t.Parallel()
+
+	srv, paths, bodies := recordingServer(t, `{"predictions":[
+		{"embeddings":{"values":[0.1,0.2],"statistics":{"token_count":3,"truncated":false}}},
+		{"embeddings":{"values":[0.3,0.4],"statistics":{"token_count":4,"truncated":false}}}
+	],"metadata":{"billableCharacterCount":11}}`)
+	client := newVertexEmbedClient(t, "gemini-embedding-001", srv.URL, map[string]any{"output_dimensionality": 2})
+
+	result, err := client.CreateBatchEmbedding(t.Context(), []string{"first", "second"})
+	require.NoError(t, err)
+
+	require.Equal(t, []string{"/v1beta1/publishers/google/models/gemini-embedding-001:predict"}, *paths)
+	body := (*bodies)[0]
+	assert.Equal(t, []any{map[string]any{"content": "first"}, map[string]any{"content": "second"}}, body["instances"])
+	assert.Equal(t, map[string]any{"outputDimensionality": float64(2)}, body["parameters"])
+
+	require.Len(t, result.Embeddings, 2)
+	assert.InDeltaSlice(t, []float64{0.1, 0.2}, result.Embeddings[0], 1e-6)
+	assert.InDeltaSlice(t, []float64{0.3, 0.4}, result.Embeddings[1], 1e-6)
+	assert.Equal(t, int64(7), result.InputTokens)
+	assert.Equal(t, int64(7), result.TotalTokens)
+}
+
+func TestCreateBatchEmbedding_VertexEmbedding2SequentialEmbedContent(t *testing.T) {
+	t.Parallel()
+
+	srv, paths, bodies := recordingServer(t,
+		`{"embedding":{"values":[0.1,0.2]},"usageMetadata":{"promptTokenCount":3},"truncated":false}`,
+		`{"embedding":{"values":[0.3,0.4]},"usageMetadata":{"promptTokenCount":4},"truncated":false}`,
+	)
+	client := newVertexEmbedClient(t, "gemini-embedding-2", srv.URL, map[string]any{"output_dimensionality": 2})
+
+	result, err := client.CreateBatchEmbedding(t.Context(), []string{"first", "second"})
+	require.NoError(t, err)
+
+	wantPath := "/v1beta1/publishers/google/models/gemini-embedding-2:embedContent"
+	require.Equal(t, []string{wantPath, wantPath}, *paths, "one embedContent request per input, in order")
+	for i, text := range []string{"first", "second"} {
+		body := (*bodies)[i]
+		assert.Equal(t, map[string]any{"role": "user", "parts": []any{map[string]any{"text": text}}}, body["content"])
+		assert.Equal(t, map[string]any{"outputDimensionality": float64(2)}, body["embedContentConfig"])
+		assert.NotContains(t, body, "instances")
+	}
+
+	require.Len(t, result.Embeddings, 2)
+	assert.InDeltaSlice(t, []float64{0.1, 0.2}, result.Embeddings[0], 1e-6)
+	assert.InDeltaSlice(t, []float64{0.3, 0.4}, result.Embeddings[1], 1e-6)
+	assert.Equal(t, int64(7), result.InputTokens)
+	assert.Equal(t, int64(7), result.TotalTokens)
+}
+
+func TestCreateBatchEmbedding_429SurfacesAsStatusError(t *testing.T) {
+	t.Parallel()
+
+	srv := embedServer(t, http.StatusTooManyRequests, `{"error":{"code":429,"message":"Resource has been exhausted","status":"RESOURCE_EXHAUSTED"}}`, nil, nil)
+	client := newEmbedClient(t, srv.URL, nil)
+
+	_, err := client.CreateBatchEmbedding(t.Context(), []string{"hello"})
+	require.Error(t, err)
+
+	var se *modelerrors.StatusError
+	require.ErrorAs(t, err, &se, "error must wrap *modelerrors.StatusError so the backoff gate can arm")
+	assert.Equal(t, http.StatusTooManyRequests, se.StatusCode)
+}
+
+func TestCreateBatchEmbedding_ContextCancelled(t *testing.T) {
+	t.Parallel()
+
+	srv := embedServer(t, http.StatusOK, `{"embeddings":[{"values":[0.1]}]}`, nil, nil)
+	client := newEmbedClient(t, srv.URL, nil)
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	_, err := client.CreateBatchEmbedding(ctx, []string{"hello"})
+	require.ErrorIs(t, err, context.Canceled)
+}
+
+func TestVertexEmbedsOneContentPerRequest(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		surface string
+		model   string
+		want    bool
+	}{
+		{apiSurfaceVertexAI, "gemini-embedding-2", true},
+		{apiSurfaceVertexAI, "gemini-embedding-001", false},
+		{apiSurfaceVertexAI, "text-embedding-005", false},
+		{apiSurfaceGeminiAPI, "gemini-embedding-2", false},
+		{apiSurfaceGateway, "gemini-embedding-2", false},
+	}
+	for _, tt := range tests {
+		c := &Client{apiSurface: tt.surface}
+		c.ModelConfig.Model = tt.model
+		assert.Equal(t, tt.want, c.vertexEmbedsOneContentPerRequest(), "%s/%s", tt.surface, tt.model)
+	}
+}
+
+func TestEmbedEach_OneVectorPerResponse(t *testing.T) {
+	t.Parallel()
+
+	var calls atomic.Int32
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var req batchEmbedRequest
+		assert.NoError(t, json.NewDecoder(r.Body).Decode(&req))
+		assert.Len(t, req.Requests, 1, "one content per request")
+		w.Header().Set("Content-Type", "application/json")
+		fmt.Fprintf(w, `{"embeddings":[{"values":[%d]}]}`, calls.Add(1))
+	}))
+	t.Cleanup(srv.Close)
+	client := newEmbedClient(t, srv.URL, nil)
+
+	genaiClient, err := client.clientFn(t.Context())
+	require.NoError(t, err)
+
+	embeddings, err := client.embedEach(t.Context(), genaiClient, []*genai.Content{
+		genai.NewContentFromText("a", genai.RoleUser),
+		genai.NewContentFromText("bb", genai.RoleUser),
+		genai.NewContentFromText("ccc", genai.RoleUser),
+	}, &genai.EmbedContentConfig{})
+	require.NoError(t, err)
+	assert.Equal(t, int32(3), calls.Load())
+	require.Len(t, embeddings, 3)
+	assert.Equal(t, []float32{1}, embeddings[0].Values)
+	assert.Equal(t, []float32{2}, embeddings[1].Values)
+	assert.Equal(t, []float32{3}, embeddings[2].Values)
+}
+
+// Two vectors then zero would still total two: the per-response check must
+// catch the misalignment instead of silently pairing texts with wrong vectors.
+func TestEmbedEach_RejectsMisalignedResponses(t *testing.T) {
+	t.Parallel()
+
+	srv, paths, _ := recordingServer(t, `{"embeddings":[{"values":[1]},{"values":[2]}]}`, `{"embeddings":[]}`)
+	client := newEmbedClient(t, srv.URL, nil)
+
+	genaiClient, err := client.clientFn(t.Context())
+	require.NoError(t, err)
+
+	_, err = client.embedEach(t.Context(), genaiClient, []*genai.Content{
+		genai.NewContentFromText("a", genai.RoleUser),
+		genai.NewContentFromText("b", genai.RoleUser),
+	}, &genai.EmbedContentConfig{})
+	require.ErrorContains(t, err, "expected 1 embedding for input 0, got 2")
+	assert.Len(t, *paths, 1, "must fail at the offending response")
+}