feat: add openai to gcp embedding translations for gemini embedding 2

internal/apischema/gcp/gcp.go

@@ -107,3 +107,47 @@ type Prediction struct {
 type PredictResponse struct {
 	Predictions []*Prediction `json:"predictions"`
 }
+
+// EmbedContentRequest is the request body for the embedContent endpoint used by newer embedding models
+// (e.g. gemini-embedding-2-*).
+// All input texts are packed as parts in a single Content object and we drop deprecated top-level fields
+// (taskType, outputDimensionality, etc.)
+//
+// See https://docs.cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.publishers.models/embedContent
+type EmbedContentRequest struct {
+	Content genai.Content       `json:"content"`
+	Config  *EmbedContentConfig `json:"embedContentConfig,omitempty"`
+}
+
+// EmbedContentConfig contains optional parameters for the embedContent method.
+// https://github.com/googleapis/go-genai/blob/v1.54.0/models.go#L727
+type EmbedContentConfig struct {
+	TaskType             openai.EmbeddingTaskType `json:"taskType,omitempty"`
+	Title                string                   `json:"title,omitempty"`
+	OutputDimensionality int                      `json:"outputDimensionality,omitempty"`
+	AutoTruncate         *bool                    `json:"autoTruncate,omitempty"`
+}
+
+// EmbedContentResponse is the response from the embedContent endpoint.
+// The REST API returns a single embedding (not an array), plus usage metadata.
+// https://docs.cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.publishers.models/embedContent
+type EmbedContentResponse struct {
+	// The embedding generated from the input content (singular object, not an array).
+	Embedding *EmbedContentEmbedding `json:"embedding,omitempty"`
+	// Usage metadata about the response.
+	UsageMetadata *EmbedContentUsageMetadata `json:"usageMetadata,omitempty"`
+	// Whether the input content was truncated before generating the embedding.
+	Truncated bool `json:"truncated,omitempty"`
+}
+
+// EmbedContentEmbedding represents the embedding values from the embedContent response.
+type EmbedContentEmbedding struct {
+	// The embedding values.
+	Values []float32 `json:"values,omitempty"`
+}
+
+// EmbedContentUsageMetadata contains token usage from the embedContent response.
+type EmbedContentUsageMetadata struct {
+	PromptTokenCount int `json:"promptTokenCount,omitempty"`
+	TotalTokenCount  int `json:"totalTokenCount,omitempty"`
+}

internal/apischema/openai/openai.go

@@ -1684,15 +1684,8 @@ type Model struct {
 	OwnedBy string `json:"owned_by"`
 }
 
-// EmbeddingRequest represents a request structure for embeddings API.
-type EmbeddingRequest struct {
-	// Input: Input text to embed, encoded as a string or array of tokens.
-	// To embed multiple inputs in a single request, pass an array of strings or array of token arrays.
-	// The input must not exceed the max input tokens for the model (8192 tokens for text-embedding-ada-002),
-	// cannot be an empty string, and any array must be 2048 dimensions or less.
-	// Docs: https://platform.openai.com/docs/api-reference/embeddings/create#embeddings-create-input
-	Input EmbeddingRequestInput `json:"input"`
-
+// EmbeddingBaseRequest holds fields shared by both embedding request variants.
+type EmbeddingBaseRequest struct {
 	// Model: ID of the model to use.
 	// Docs: https://platform.openai.com/docs/api-reference/embeddings/create#embeddings-create-model
 	Model string `json:"model"`
@@ -1714,6 +1707,76 @@ type EmbeddingRequest struct {
 	*GCPVertexAIEmbeddingVendorFields `json:",inline,omitempty"`
 }
 
+// EmbeddingCompletionRequest is the text-only embedding request (classic OpenAI style).
+// https://developers.openai.com/api/reference/resources/embeddings/methods/create
+type EmbeddingCompletionRequest struct {
+	EmbeddingBaseRequest
+	// Input: Input text to embed, encoded as a string or array of tokens.
+	// To embed multiple inputs in a single request, pass an array of strings or array of token arrays.
+	// The input must not exceed the max input tokens for the model (8192 tokens for text-embedding-ada-002),
+	// cannot be an empty string, and any array must be 2048 dimensions or less.
+	Input EmbeddingRequestInput `json:"input"`
+}
+
+// EmbeddingChatRequest is the chat-style embedding request (supports multimodal).
+// Following vLLM convention: https://github.com/vllm-project/vllm/blob/2a16ece2d342c0c154a4949ad317b521f8c04ec4/vllm/entrypoints/pooling/embed/protocol.py#L83
+type EmbeddingChatRequest struct {
+	EmbeddingBaseRequest
+	// Messages: Chat messages for multimodal embedding.
+	// Uses the same chat message format as /v1/chat/completions.
+	// Only user messages are processed; system/assistant/tool messages are ignored.
+	Messages []ChatCompletionMessageParamUnion `json:"messages"`
+}
+
+// EmbeddingRequest is a discriminated union: exactly one of OfCompletion or OfChat is set.
+// Discrimination is by presence of "input" (→ completion) vs "messages" (→ chat) in JSON.
+type EmbeddingRequest struct {
+	EmbeddingBaseRequest                             // promoted shared fields
+	OfCompletion         *EmbeddingCompletionRequest // set when JSON has "input"
+	OfChat               *EmbeddingChatRequest       // set when JSON has "messages"
+}
+
+// UnmarshalJSON discriminates between completion and chat embedding requests.
+func (r *EmbeddingRequest) UnmarshalJSON(data []byte) error {
+	hasInput := gjson.GetBytes(data, "input").Exists()
+	hasMessages := gjson.GetBytes(data, "messages").Exists()
+
+	if hasInput && hasMessages {
+		return fmt.Errorf("embedding request must have either 'input' or 'messages', not both")
+	}
+
+	if hasMessages {
+		var chat EmbeddingChatRequest
+		if err := json.Unmarshal(data, &chat); err != nil {
+			return err
+		}
+		r.EmbeddingBaseRequest = chat.EmbeddingBaseRequest
+		r.OfChat = &chat
+		return nil
+	}
+
+	// Default to completion (input-based) — this handles both explicit "input" and legacy cases.
+	var comp EmbeddingCompletionRequest
+	if err := json.Unmarshal(data, &comp); err != nil {
+		return err
+	}
+	r.EmbeddingBaseRequest = comp.EmbeddingBaseRequest
+	r.OfCompletion = &comp
+	return nil
+}
+
+// MarshalJSON delegates to the active variant.
+func (r EmbeddingRequest) MarshalJSON() ([]byte, error) {
+	if r.OfChat != nil {
+		return json.Marshal(r.OfChat)
+	}
+	if r.OfCompletion != nil {
+		return json.Marshal(r.OfCompletion)
+	}
+	// Fallback: marshal just the base fields.
+	return json.Marshal(r.EmbeddingBaseRequest)
+}
+
 type EmbeddingTaskType string
 
 const (
@@ -1727,16 +1790,20 @@ const (
 	EmbeddingTaskTypeCodeRetrievalQuery EmbeddingTaskType = "CODE_RETRIEVAL_QUERY"
 )
 
-// GCPVertexAIEmbeddingVendorFields contains GCP Vertex AI (Gemini) vendor-specific fields for embeddings.
+// GCPVertexAIEmbeddingVendorFields contains GCP Vertex AI vendor-specific fields for embeddings.
+// The translator maps these to the appropriate wire format per endpoint:
+//   - predict: parameters.auto_truncate, instances[].task_type, instances[].title
+//   - embedContent: embedContentConfig.autoTruncate, embedContentConfig.taskType, embedContentConfig.title
 type GCPVertexAIEmbeddingVendorFields struct {
-	// When set to true, input text will be truncated. When set to false, an error is returned if the input text is longer than the maximum length supported by the model. Defaults to true.
-	// https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api#parameter-list
+	// Auto-truncate input text if it exceeds the model's max length. Defaults to true.
+	AutoTruncate *bool `json:"auto_truncate,omitempty"`
 
-	AutoTruncate bool `json:"auto_truncate,omitempty"`
-
-	// This is global task_type set, which is convenient for users. If left blank, the default used is RETRIEVAL_QUERY.
-	// For more information about task types, see https://docs.cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types
+	// Global task type for the request. Defaults to RETRIEVAL_QUERY if unset.
+	// https://docs.cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types
 	TaskType EmbeddingTaskType `json:"task_type,omitempty"`
+
+	// Title for the embedding content. Helps the model produce better embeddings.
+	Title string `json:"title,omitempty"`
 }
 
 // EmbeddingResponse represents a response from /v1/embeddings.

internal/endpointspec/endpointspec_test.go

@@ -145,8 +145,14 @@ func TestEmbeddingsEndpointSpec_ParseBody(t *testing.T) {
 		require.ErrorContains(t, err, "malformed request")
 	})
 
-	t.Run("success", func(t *testing.T) {
-		req := openai.EmbeddingRequest{Model: "text-embedding-3-large", Input: openai.EmbeddingRequestInput{Value: "input"}}
+	t.Run("success with input", func(t *testing.T) {
+		req := openai.EmbeddingRequest{
+			EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-large"},
+			OfCompletion: &openai.EmbeddingCompletionRequest{
+				EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-large"},
+				Input:                openai.EmbeddingRequestInput{Value: "input"},
+			},
+		}
 		body, err := json.Marshal(req)
 		require.NoError(t, err)
 
@@ -157,6 +163,17 @@ func TestEmbeddingsEndpointSpec_ParseBody(t *testing.T) {
 		require.NotNil(t, parsed)
 		require.Nil(t, mutated)
 	})
+
+	t.Run("success with messages", func(t *testing.T) {
+		body := []byte(`{"model":"gemini-embedding-2","messages":[{"role":"user","content":"embed this"}]}`)
+
+		model, parsed, stream, mutated, err := spec.ParseBody(body, false)
+		require.NoError(t, err)
+		require.Equal(t, "gemini-embedding-2", model)
+		require.False(t, stream)
+		require.NotNil(t, parsed)
+		require.Nil(t, mutated)
+	})
 }
 
 func TestEmbeddingsEndpointSpec_GetTranslator(t *testing.T) {

internal/tracing/openinference/openai/embeddings_test.go

@@ -24,23 +24,46 @@ import (
 // Test data.
 var (
 	basicEmbeddingReq = &openai.EmbeddingRequest{
-		Model: "text-embedding-3-small",
-		Input: openai.EmbeddingRequestInput{Value: "How do I reset my password?"},
+		EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
+		OfCompletion: &openai.EmbeddingCompletionRequest{
+			EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
+			Input:                openai.EmbeddingRequestInput{Value: "How do I reset my password?"},
+		},
 	}
 	basicEmbeddingReqBody = []byte(`{"model":"text-embedding-3-small","input":"How do I reset my password?"}`)
 
 	multiInputEmbeddingReq = &openai.EmbeddingRequest{
-		Model: "text-embedding-3-small",
-		Input: openai.EmbeddingRequestInput{Value: []string{"How", "do", "I", "reset", "my", "password?"}},
+		EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
+		OfCompletion: &openai.EmbeddingCompletionRequest{
+			EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
+			Input:                openai.EmbeddingRequestInput{Value: []string{"How", "do", "I", "reset", "my", "password?"}},
+		},
 	}
 	multiInputEmbeddingReqBody = []byte(`{"model":"text-embedding-3-small","input":["How","do","I","reset","my","password?"]}`)
 
 	tokenInputEmbeddingReq = &openai.EmbeddingRequest{
-		Model: "text-embedding-3-small",
-		Input: openai.EmbeddingRequestInput{Value: []int{14438, 656, 358, 7738, 856, 3636, 30}},
+		EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
+		OfCompletion: &openai.EmbeddingCompletionRequest{
+			EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
+			Input:                openai.EmbeddingRequestInput{Value: []int64{4438, 656, 358, 7738, 856, 3636, 30}},
+		},
 	}
 	tokenInputEmbeddingReqBody = []byte(`{"model":"text-embedding-3-small","input":[4438,656,358,7738,856,3636,30]}`)
 
+	chatEmbeddingReq = &openai.EmbeddingRequest{
+		EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "gemini-embedding-2"},
+		OfChat: &openai.EmbeddingChatRequest{
+			EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "gemini-embedding-2"},
+			Messages: []openai.ChatCompletionMessageParamUnion{
+				{OfUser: &openai.ChatCompletionUserMessageParam{
+					Role:    "user",
+					Content: openai.StringOrUserRoleContentUnion{Value: "embed this text"},
+				}},
+			},
+		},
+	}
+	chatEmbeddingReqBody = []byte(`{"model":"gemini-embedding-2","messages":[{"role":"user","content":"embed this text"}]}`)
+
 	basicEmbeddingResp = &openai.EmbeddingResponse{
 		Model: "text-embedding-3-small",
 		Usage: openai.EmbeddingUsage{
@@ -95,6 +118,12 @@ func TestEmbeddingsRecorder_StartParams(t *testing.T) {
 			reqBody:          multiInputEmbeddingReqBody,
 			expectedSpanName: "CreateEmbeddings",
 		},
+		{
+			name:             "chat embedding request",
+			req:              chatEmbeddingReq,
+			reqBody:          chatEmbeddingReqBody,
+			expectedSpanName: "CreateEmbeddings",
+		},
 	}
 
 	for _, tt := range tests {
@@ -172,6 +201,19 @@ func TestEmbeddingsRecorder_RecordRequest(t *testing.T) {
 				attribute.String(openinference.EmbeddingInvocationParameters, `{"model":"text-embedding-3-small"}`),
 			},
 		},
+		{
+			name:    "chat embedding request",
+			req:     chatEmbeddingReq,
+			reqBody: chatEmbeddingReqBody,
+			config:  &openinference.TraceConfig{},
+			expectedAttrs: []attribute.KeyValue{
+				attribute.String(openinference.SpanKind, openinference.SpanKindEmbedding),
+				attribute.String(openinference.InputValue, string(chatEmbeddingReqBody)),
+				attribute.String(openinference.InputMimeType, openinference.MimeTypeJSON),
+				attribute.String(openinference.EmbeddingInvocationParameters, `{"model":"gemini-embedding-2"}`),
+				attribute.String(openinference.EmbeddingTextAttribute(0), "embed this text"),
+			},
+		},
 		{
 			name:    "hidden invocation parameters",
 			req:     basicEmbeddingReq,

internal/tracing/openinference/openai/request_attrs.go

@@ -256,16 +256,40 @@ func buildEmbeddingsRequestAttributes(embRequest *openai.EmbeddingRequest, body
 	// 4. Azure deployments don't affect this (they only host OpenAI models with cl100k_base)
 	// Following OpenInference spec guidance to only record human-readable text.
 	if !config.HideInputs && !config.HideEmbeddingsText {
-		switch input := embRequest.Input.Value.(type) {
-		case string:
-			attrs = append(attrs, attribute.String(openinference.EmbeddingTextAttribute(0), input))
-		case []string:
-			for i, text := range input {
-				attrs = append(attrs, attribute.String(openinference.EmbeddingTextAttribute(i), text))
+		switch {
+		case embRequest.OfCompletion != nil:
+			switch input := embRequest.OfCompletion.Input.Value.(type) {
+			case string:
+				attrs = append(attrs, attribute.String(openinference.EmbeddingTextAttribute(0), input))
+			case []string:
+				for i, text := range input {
+					attrs = append(attrs, attribute.String(openinference.EmbeddingTextAttribute(i), text))
+				}
+			// Token inputs are not recorded to reduce span size.
+			case []int64:
+			case [][]int64:
+			}
+		case embRequest.OfChat != nil:
+			idx := 0
+			for _, msg := range embRequest.OfChat.Messages {
+				if msg.OfUser == nil {
+					continue
+				}
+				switch content := msg.OfUser.Content.Value.(type) {
+				case string:
+					if content != "" {
+						attrs = append(attrs, attribute.String(openinference.EmbeddingTextAttribute(idx), content))
+						idx++
+					}
+				case []openai.ChatCompletionContentPartUserUnionParam:
+					for _, part := range content {
+						if part.OfText != nil && part.OfText.Text != "" {
+							attrs = append(attrs, attribute.String(openinference.EmbeddingTextAttribute(idx), part.OfText.Text))
+							idx++
+						}
+					}
+				}
 			}
-		// Token inputs are not recorded to reduce span size.
-		case []int64:
-		case [][]int64:
 		}
 	}
 

internal/tracing/tracing_test.go

@@ -775,9 +775,12 @@ func TestNewTracingFromEnv_Embeddings_Redaction(t *testing.T) {
 
 			// Create a test request with sensitive data.
 			req := &openai.EmbeddingRequest{
-				Model: "text-embedding-3-small",
-				Input: openai.EmbeddingRequestInput{
-					Value: "Sensitive embedding text",
+				EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
+				OfCompletion: &openai.EmbeddingCompletionRequest{
+					EmbeddingBaseRequest: openai.EmbeddingBaseRequest{Model: "text-embedding-3-small"},
+					Input: openai.EmbeddingRequestInput{
+						Value: "Sensitive embedding text",
+					},
 				},
 			}
 			reqBody := []byte(`{"input":"Sensitive embedding text","model":"text-embedding-3-small"}`)

internal/translator/openai_awsbedrock_embeddings.go

@@ -41,8 +41,12 @@ func (o *openAIToAWSBedrockTranslatorV1Embedding) RequestBody(_ []byte, req *ope
 	}
 	o.requestModel = model
 
+	if req.OfCompletion == nil {
+		return nil, nil, fmt.Errorf("%w: AWS Bedrock Titan requires an input-based embedding request (messages not supported)", internalapi.ErrInvalidRequestBody)
+	}
+
 	var inputText string
-	switch v := req.Input.Value.(type) {
+	switch v := req.OfCompletion.Input.Value.(type) {
 	case string:
 		inputText = v
 	case []string:
@@ -52,7 +56,7 @@ func (o *openAIToAWSBedrockTranslatorV1Embedding) RequestBody(_ []byte, req *ope
 		}
 		inputText = v[0]
 	default:
-		return nil, nil, fmt.Errorf("%w: unsupported input type %T", internalapi.ErrInvalidRequestBody, req.Input.Value)
+		return nil, nil, fmt.Errorf("%w: unsupported input type %T", internalapi.ErrInvalidRequestBody, req.OfCompletion.Input.Value)
 	}
 
 	bedrockReq := awsbedrock.TitanEmbeddingRequest{