update

Signed-off-by: yxia216 <[email protected]>

Author: hustxiayang

internal/apischema/gcp/gcp.go

@@ -38,11 +38,13 @@ type GenerateContentRequest struct {
 }
 
 type EmbedContentRequest struct {
-	// Contains the multipart content of a message.
-	//
-	// https://github.com/googleapis/go-genai/blob/6a8184fcaf8bf15f0c566616a7b356560309be9b/types.go#L858
-	Contents []genai.Content `json:"contents"`
-	// Tool details of a tool that the model may use to generate a response.
+	// Content to be embedded. Only text content is supported for embeddings.
+	Content *genai.Content `json:"content"`
 
+	// Optional configuration for the embedding request.
+	// Uses the official genai library configuration structure.
 	Config *genai.EmbedContentConfig `json:"config,omitempty"`
 }
+
+// Note: We now use genai.EmbedContentResponse directly instead of defining our own.
+// This provides better compatibility and includes metadata like token usage.

internal/apischema/openai/openai.go

@@ -1523,6 +1523,9 @@ type EmbeddingCompletionRequest struct {
 	// User: A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.
 	// Docs: https://platform.openai.com/docs/api-reference/embeddings/create#embeddings-create-user
 	User *string `json:"user,omitempty"`
+
+	// GCPVertexAIEmbeddingVendorFields configures the GCP VertexAI specific fields during schema translation.
+	*GCPVertexAIEmbeddingVendorFields `json:",inline,omitempty"`
 }
 
 // GetModel implements ModelName interface
@@ -1553,26 +1556,67 @@ type EmbeddingChatRequest struct {
 	// User: A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.
 	// Docs: https://platform.openai.com/docs/api-reference/embeddings/create#embeddings-create-user
 	User *string `json:"user,omitempty"`
+
+	// GCPVertexAIEmbeddingVendorFields configures the GCP VertexAI specific fields during schema translation.
+	*GCPVertexAIEmbeddingVendorFields `json:",inline,omitempty"`
 }
 
 // GetModel implements ModelProvider interface
 func (e *EmbeddingChatRequest) GetModel() string {
 	return e.Model
 }
 
-type EmbeddingRequest interface {
-	EmbeddingCompletionRequest | EmbeddingChatRequest
+// EmbeddingRequest is a union type that can handle both EmbeddingCompletionRequest and EmbeddingChatRequest.
+type EmbeddingRequest struct {
+	OfCompletion *EmbeddingCompletionRequest `json:",omitzero,inline"`
+	OfChat       *EmbeddingChatRequest       `json:",omitzero,inline"`
+}
+
+// UnmarshalJSON implements json.Unmarshaler to handle both EmbeddingCompletionRequest and EmbeddingChatRequest.
+func (e *EmbeddingRequest) UnmarshalJSON(data []byte) error {
+	// Check for Messages field to distinguish EmbeddingChatRequest
+	messagesResult := gjson.GetBytes(data, "messages")
+	if messagesResult.Exists() {
+		var chatReq EmbeddingChatRequest
+		if err := json.Unmarshal(data, &chatReq); err != nil {
+			return err
+		}
+		e.OfChat = &chatReq
+		return nil
+	}
+
+	// Check for Input field to distinguish EmbeddingCompletionRequest
+	inputResult := gjson.GetBytes(data, "input")
+	if inputResult.Exists() {
+		var completionReq EmbeddingCompletionRequest
+		if err := json.Unmarshal(data, &completionReq); err != nil {
+			return err
+		}
+		e.OfCompletion = &completionReq
+		return nil
+	}
+
+	return errors.New("embedding request must have either 'input' field (EmbeddingCompletionRequest) or 'messages' field (EmbeddingChatRequest)")
 }
 
-// ModelName interface for types that can provide a model name
-type ModelName interface {
-	GetModel() string
+// MarshalJSON implements json.Marshaler.
+func (e EmbeddingRequest) MarshalJSON() ([]byte, error) {
+	if e.OfCompletion != nil {
+		return json.Marshal(e.OfCompletion)
+	}
+	if e.OfChat != nil {
+		return json.Marshal(e.OfChat)
+	}
+	return nil, errors.New("no embedding request to marshal")
 }
 
 // GetModelFromEmbeddingRequest extracts the model name from any EmbeddingRequest type
-func GetModelFromEmbeddingRequest[T EmbeddingRequest](req *T) string {
-	if mp, ok := any(*req).(ModelName); ok {
-		return mp.GetModel()
+func GetModelFromEmbeddingRequest(req *EmbeddingRequest) string {
+	if req.OfCompletion != nil {
+		return req.OfCompletion.GetModel()
+	}
+	if req.OfChat != nil {
+		return req.OfChat.GetModel()
 	}
 	return ""
 }
@@ -1650,6 +1694,13 @@ type EmbeddingUsage struct {
 	TotalTokens int `json:"total_tokens"` //nolint:tagliatelle //follow openai api
 }
 
+// GCPVertexAIEmbeddingVendorFields contains GCP Vertex AI (Gemini) vendor-specific fields for embedding requests.
+type GCPVertexAIEmbeddingVendorFields struct {
+	// Type of task for which the embedding will be used.
+	// https://docs.cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types#supported_task_types
+	TaskType string `json:"task_type,omitempty"`
+}
+
 // JSONUNIXTime is a helper type to marshal/unmarshal time.Time UNIX timestamps.
 type JSONUNIXTime time.Time
 

internal/tracing/openinference/openai/request_attrs.go

@@ -231,11 +231,29 @@ func buildEmbeddingsRequestAttributes(embRequest *openai.EmbeddingRequest, body
 	}
 
 	if !config.HideLLMInvocationParameters {
+		// Extract parameters from the union type
+		var model string
+		var encodingFormat *string
+		var dimensions *int
+		var user *string
+
+		if embRequest.OfCompletion != nil {
+			model = embRequest.OfCompletion.Model
+			encodingFormat = embRequest.OfCompletion.EncodingFormat
+			dimensions = embRequest.OfCompletion.Dimensions
+			user = embRequest.OfCompletion.User
+		} else if embRequest.OfChat != nil {
+			model = embRequest.OfChat.Model
+			encodingFormat = embRequest.OfChat.EncodingFormat
+			dimensions = embRequest.OfChat.Dimensions
+			user = embRequest.OfChat.User
+		}
+
 		params := embeddingsInvocationParameters{
-			Model:          embRequest.Model,
-			EncodingFormat: embRequest.EncodingFormat,
-			Dimensions:     embRequest.Dimensions,
-			User:           embRequest.User,
+			Model:          model,
+			EncodingFormat: encodingFormat,
+			Dimensions:     dimensions,
+			User:           user,
 		}
 		if invocationParamsJSON, err := json.Marshal(params); err == nil {
 			attrs = append(attrs, attribute.String(openinference.EmbeddingInvocationParameters, string(invocationParamsJSON)))
@@ -250,16 +268,26 @@ 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))
+		var inputValue any
+		if embRequest.OfCompletion != nil {
+			inputValue = embRequest.OfCompletion.Input.Value
+		} else if embRequest.OfChat != nil {
+			// For chat requests, we'll extract text from messages
+			inputValue = "chat_messages" // Simplified - could be enhanced to extract actual text
+		}
+
+		if inputValue != nil {
+			switch input := inputValue.(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:
 			}
-		// Token inputs are not recorded to reduce span size.
-		case []int64:
-		case [][]int64:
 		}
 	}
 

internal/translator/openai_azureopenai_embeddings.go

@@ -33,10 +33,10 @@ type openAIToAzureOpenAITranslatorV1Embedding struct {
 }
 
 // RequestBody implements [OpenAIEmbeddingTranslator.RequestBody].
-func (o *openAIToAzureOpenAITranslatorV1Embedding) RequestBody(original []byte, req *openai.EmbeddingCompletionRequest, onRetry bool) (
+func (o *openAIToAzureOpenAITranslatorV1Embedding) RequestBody(original []byte, req *openai.EmbeddingRequest, onRetry bool) (
 	newHeaders []internalapi.Header, newBody []byte, err error,
 ) {
-	modelName := req.Model
+	modelName := openai.GetModelFromEmbeddingRequest(req)
 	if o.modelNameOverride != "" {
 		// If modelName is set we override the model to be used for the request.
 		newBody, err = sjson.SetBytesOptions(original, "model", o.modelNameOverride, sjsonOptions)

internal/translator/openai_embeddings.go

@@ -36,7 +36,7 @@ type openAIToOpenAITranslatorV1Embedding struct {
 }
 
 // RequestBody implements [OpenAIEmbeddingTranslator.RequestBody].
-func (o *openAIToOpenAITranslatorV1Embedding) RequestBody(original []byte, _ *openai.EmbeddingCompletionRequest, onRetry bool) (
+func (o *openAIToOpenAITranslatorV1Embedding) RequestBody(original []byte, req *openai.EmbeddingRequest, onRetry bool) (
 	newHeaders []internalapi.Header, newBody []byte, err error,
 ) {
 	if o.modelNameOverride != "" {

internal/translator/openai_gcpvertexai.go

@@ -533,10 +533,9 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) geminiResponseToOpenAIMe
 	return openaiResp, nil
 }
 
-// ResponseError implements [OpenAIChatCompletionTranslator.ResponseError].
-// Translate GCP Vertex AI exceptions to OpenAI error type.
-// GCP error responses typically contain JSON with error details or plain text error messages.
-func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) ResponseError(respHeaders map[string]string, body io.Reader) (
+// convertGCPVertexAIErrorToOpenAI converts GCP Vertex AI error responses to OpenAI error format.
+// This is a shared function used by both chat completion and embedding translators.
+func convertGCPVertexAIErrorToOpenAI(respHeaders map[string]string, body io.Reader) (
 	newHeaders []internalapi.Header, newBody []byte, err error,
 ) {
 	var buf []byte
@@ -545,8 +544,8 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) ResponseError(respHeader
 		return nil, nil, fmt.Errorf("failed to read error body: %w", err)
 	}
 
-	// Assume all responses have a valid status code header.
 	statusCode := respHeaders[statusHeaderName]
+	contentType := respHeaders[contentTypeHeaderName]
 
 	openaiError := openai.Error{
 		Type: "error",
@@ -556,19 +555,45 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) ResponseError(respHeader
 		},
 	}
 
-	var gcpError gcpVertexAIError
-	// Try to parse as GCP error response structure.
-	if err = json.Unmarshal(buf, &gcpError); err == nil {
-		errMsg := gcpError.Error.Message
-		if len(gcpError.Error.Details) > 0 {
-			// If details are present and not null, append them to the error message.
-			errMsg = fmt.Sprintf("Error: %s\nDetails: %s", errMsg, string(gcpError.Error.Details))
-		}
-		openaiError.Error.Type = gcpError.Error.Status
-		openaiError.Error.Message = errMsg
-	} else {
-		// If not JSON, read the raw body as the error message.
+	// If the content type is not JSON, treat it as a generic error
+	if contentType != "" && contentType != jsonContentType {
 		openaiError.Error.Message = string(buf)
+	} else {
+		var gcpError gcpVertexAIError
+		// Try to parse as GCP error response structure first
+		if err = json.Unmarshal(buf, &gcpError); err == nil {
+			errMsg := gcpError.Error.Message
+			if len(gcpError.Error.Details) > 0 {
+				// If details are present and not null, append them to the error message.
+				errMsg = fmt.Sprintf("Error: %s\nDetails: %s", errMsg, string(gcpError.Error.Details))
+			}
+			openaiError.Error.Type = gcpError.Error.Status
+			openaiError.Error.Message = errMsg
+		} else {
+			// Try to parse as generic JSON error format
+			var genericError map[string]interface{}
+			if err := json.Unmarshal(buf, &genericError); err == nil {
+				// Extract error message from generic JSON error format
+				var errorMessage string
+				if errorField, exists := genericError["error"]; exists {
+					if errorMap, ok := errorField.(map[string]interface{}); ok {
+						if message, exists := errorMap["message"]; exists {
+							if msgStr, ok := message.(string); ok {
+								errorMessage = msgStr
+							}
+						}
+					}
+				}
+				if errorMessage != "" {
+					openaiError.Error.Message = errorMessage
+				} else {
+					openaiError.Error.Message = string(buf)
+				}
+			} else {
+				// If not parseable as JSON, use raw body as the error message
+				openaiError.Error.Message = string(buf)
+			}
+		}
 	}
 
 	newBody, err = json.Marshal(openaiError)
@@ -581,3 +606,12 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) ResponseError(respHeader
 	}
 	return
 }
+
+// ResponseError implements [OpenAIChatCompletionTranslator.ResponseError].
+// Translate GCP Vertex AI exceptions to OpenAI error type.
+// GCP error responses typically contain JSON with error details or plain text error messages.
+func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) ResponseError(respHeaders map[string]string, body io.Reader) (
+	newHeaders []internalapi.Header, newBody []byte, err error,
+) {
+	return convertGCPVertexAIErrorToOpenAI(respHeaders, body)
+}