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

@@ -1518,6 +1518,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
@@ -1548,26 +1551,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 ""
 }
@@ -1645,6 +1689,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/extproc/embeddings_processor.go

@@ -32,14 +32,14 @@ func EmbeddingsProcessorFactory(f metrics.Factory) ProcessorFactory {
 	return func(config *filterapi.RuntimeConfig, requestHeaders map[string]string, logger *slog.Logger, tracing tracing.Tracing, isUpstreamFilter bool) (Processor, error) {
 		logger = logger.With("processor", "embeddings", "isUpstreamFilter", fmt.Sprintf("%v", isUpstreamFilter))
 		if !isUpstreamFilter {
-			return &embeddingsProcessorRouterFilter[openai.EmbeddingCompletionRequest]{
+			return &embeddingsProcessorRouterFilter{
 				config:         config,
 				tracer:         tracing.EmbeddingsTracer(),
 				requestHeaders: requestHeaders,
 				logger:         logger,
 			}, nil
 		}
-		return &embeddingsProcessorUpstreamFilter[openai.EmbeddingCompletionRequest]{
+		return &embeddingsProcessorUpstreamFilter{
 			config:         config,
 			requestHeaders: requestHeaders,
 			logger:         logger,
@@ -51,7 +51,7 @@ func EmbeddingsProcessorFactory(f metrics.Factory) ProcessorFactory {
 // embeddingsProcessorRouterFilter implements [Processor] for the `/v1/embeddings` endpoint.
 //
 // This is primarily used to select the route for the request based on the model name.
-type embeddingsProcessorRouterFilter[T openai.EmbeddingRequest] struct {
+type embeddingsProcessorRouterFilter struct {
 	passThroughProcessor
 	// upstreamFilter is the upstream filter that is used to process the request at the upstream filter.
 	// This will be updated when the request is retried.
@@ -67,7 +67,7 @@ type embeddingsProcessorRouterFilter[T openai.EmbeddingRequest] struct {
 	// originalRequestBody is the original request body that is passed to the upstream filter.
 	// This is used to perform the transformation of the request body on the original input
 	// when the request is retried.
-	originalRequestBody    *T
+	originalRequestBody    *openai.EmbeddingRequest
 	originalRequestBodyRaw []byte
 	// tracer is the tracer used for requests.
 	tracer tracing.EmbeddingsTracer
@@ -79,7 +79,7 @@ type embeddingsProcessorRouterFilter[T openai.EmbeddingRequest] struct {
 }
 
 // ProcessResponseHeaders implements [Processor.ProcessResponseHeaders].
-func (e *embeddingsProcessorRouterFilter[T]) ProcessResponseHeaders(ctx context.Context, headerMap *corev3.HeaderMap) (*extprocv3.ProcessingResponse, error) {
+func (e *embeddingsProcessorRouterFilter) ProcessResponseHeaders(ctx context.Context, headerMap *corev3.HeaderMap) (*extprocv3.ProcessingResponse, error) {
 	// If the request failed to route and/or immediate response was returned before the upstream filter was set,
 	// e.upstreamFilter can be nil.
 	if e.upstreamFilter != nil { // See the comment on the "upstreamFilter" field.
@@ -89,7 +89,7 @@ func (e *embeddingsProcessorRouterFilter[T]) ProcessResponseHeaders(ctx context.
 }
 
 // ProcessResponseBody implements [Processor.ProcessResponseBody].
-func (e *embeddingsProcessorRouterFilter[T]) ProcessResponseBody(ctx context.Context, body *extprocv3.HttpBody) (*extprocv3.ProcessingResponse, error) {
+func (e *embeddingsProcessorRouterFilter) ProcessResponseBody(ctx context.Context, body *extprocv3.HttpBody) (*extprocv3.ProcessingResponse, error) {
 	// If the request failed to route and/or immediate response was returned before the upstream filter was set,
 	// e.upstreamFilter can be nil.
 	if e.upstreamFilter != nil { // See the comment on the "upstreamFilter" field.
@@ -99,8 +99,8 @@ func (e *embeddingsProcessorRouterFilter[T]) ProcessResponseBody(ctx context.Con
 }
 
 // ProcessRequestBody implements [Processor.ProcessRequestBody].
-func (e *embeddingsProcessorRouterFilter[T]) ProcessRequestBody(ctx context.Context, rawBody *extprocv3.HttpBody) (*extprocv3.ProcessingResponse, error) {
-	originalModel, body, err := parseOpenAIEmbeddingBody[T](rawBody)
+func (e *embeddingsProcessorRouterFilter) ProcessRequestBody(ctx context.Context, rawBody *extprocv3.HttpBody) (*extprocv3.ProcessingResponse, error) {
+	originalModel, body, err := parseOpenAIEmbeddingBody(rawBody)
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse request body: %w", err)
 	}
@@ -125,7 +125,7 @@ func (e *embeddingsProcessorRouterFilter[T]) ProcessRequestBody(ctx context.Cont
 		ctx,
 		e.requestHeaders,
 		&headerMutationCarrier{m: headerMutation},
-		convertToEmbeddingCompletionRequest(body),
+		body,
 		rawBody.Body,
 	)
 
@@ -144,7 +144,7 @@ func (e *embeddingsProcessorRouterFilter[T]) ProcessRequestBody(ctx context.Cont
 // embeddingsProcessorUpstreamFilter implements [Processor] for the `/v1/embeddings` endpoint at the upstream filter.
 //
 // This is created per retry and handles the translation as well as the authentication of the request.
-type embeddingsProcessorUpstreamFilter[T openai.EmbeddingRequest] struct {
+type embeddingsProcessorUpstreamFilter struct {
 	logger                 *slog.Logger
 	config                 *filterapi.RuntimeConfig
 	requestHeaders         map[string]string
@@ -156,7 +156,7 @@ type embeddingsProcessorUpstreamFilter[T openai.EmbeddingRequest] struct {
 	headerMutator          *headermutator.HeaderMutator
 	bodyMutator            *bodymutator.BodyMutator
 	originalRequestBodyRaw []byte
-	originalRequestBody    *T
+	originalRequestBody    *openai.EmbeddingRequest
 	translator             translator.OpenAIEmbeddingTranslator
 	// onRetry is true if this is a retry request at the upstream filter.
 	onRetry bool
@@ -169,14 +169,14 @@ type embeddingsProcessorUpstreamFilter[T openai.EmbeddingRequest] struct {
 }
 
 // selectTranslator selects the translator based on the output schema.
-func (e *embeddingsProcessorUpstreamFilter[T]) selectTranslator(out filterapi.VersionedAPISchema) error {
+func (e *embeddingsProcessorUpstreamFilter) selectTranslator(out filterapi.VersionedAPISchema) error {
 	switch out.Name {
 	case filterapi.APISchemaOpenAI:
 		e.translator = translator.NewEmbeddingOpenAIToOpenAITranslator(out.Version, e.modelNameOverride)
 	case filterapi.APISchemaAzureOpenAI:
 		e.translator = translator.NewEmbeddingOpenAIToAzureOpenAITranslator(out.Version, e.modelNameOverride)
 	case filterapi.APISchemaGCPVertexAI:
-		e.translator = translator.NewEmbeddingOpenAIToAzureOpenAITranslator(out.Version, e.modelNameOverride)
+		e.translator = translator.NewEmbeddingOpenAIToGCPVertexAITranslator("", e.modelNameOverride)
 	default:
 		return fmt.Errorf("unsupported API schema: backend=%s", out)
 	}
@@ -189,7 +189,7 @@ func (e *embeddingsProcessorUpstreamFilter[T]) selectTranslator(out filterapi.Ve
 // So, we simply do the translation and upstream auth at this stage, and send them back to Envoy
 // with the status CONTINUE_AND_REPLACE. This will allows Envoy to not send the request body again
 // to the extproc.
-func (e *embeddingsProcessorUpstreamFilter[T]) ProcessRequestHeaders(ctx context.Context, _ *corev3.HeaderMap) (res *extprocv3.ProcessingResponse, err error) {
+func (e *embeddingsProcessorUpstreamFilter) ProcessRequestHeaders(ctx context.Context, _ *corev3.HeaderMap) (res *extprocv3.ProcessingResponse, err error) {
 	defer func() {
 		if err != nil {
 			e.metrics.RecordRequestCompletion(ctx, false, e.requestHeaders)
@@ -204,7 +204,7 @@ func (e *embeddingsProcessorUpstreamFilter[T]) ProcessRequestHeaders(ctx context
 	reqModel := cmp.Or(e.requestHeaders[internalapi.ModelNameHeaderKeyDefault], openai.GetModelFromEmbeddingRequest(e.originalRequestBody))
 	e.metrics.SetRequestModel(reqModel)
 
-	newHeaders, newBody, err := e.translator.RequestBody(e.originalRequestBodyRaw, convertToEmbeddingCompletionRequest(e.originalRequestBody), e.onRetry)
+	newHeaders, newBody, err := e.translator.RequestBody(e.originalRequestBodyRaw, e.originalRequestBody, e.onRetry)
 	if err != nil {
 		return nil, fmt.Errorf("failed to transform request: %w", err)
 	}
@@ -267,12 +267,12 @@ func (e *embeddingsProcessorUpstreamFilter[T]) ProcessRequestHeaders(ctx context
 }
 
 // ProcessRequestBody implements [Processor.ProcessRequestBody].
-func (e *embeddingsProcessorUpstreamFilter[T]) ProcessRequestBody(context.Context, *extprocv3.HttpBody) (res *extprocv3.ProcessingResponse, err error) {
+func (e *embeddingsProcessorUpstreamFilter) ProcessRequestBody(context.Context, *extprocv3.HttpBody) (res *extprocv3.ProcessingResponse, err error) {
 	panic("BUG: ProcessRequestBody should not be called in the upstream filter")
 }
 
 // ProcessResponseHeaders implements [Processor.ProcessResponseHeaders].
-func (e *embeddingsProcessorUpstreamFilter[T]) ProcessResponseHeaders(ctx context.Context, headers *corev3.HeaderMap) (res *extprocv3.ProcessingResponse, err error) {
+func (e *embeddingsProcessorUpstreamFilter) ProcessResponseHeaders(ctx context.Context, headers *corev3.HeaderMap) (res *extprocv3.ProcessingResponse, err error) {
 	defer func() {
 		if err != nil {
 			e.metrics.RecordRequestCompletion(ctx, false, e.requestHeaders)
@@ -296,7 +296,7 @@ func (e *embeddingsProcessorUpstreamFilter[T]) ProcessResponseHeaders(ctx contex
 }
 
 // ProcessResponseBody implements [Processor.ProcessResponseBody].
-func (e *embeddingsProcessorUpstreamFilter[T]) ProcessResponseBody(ctx context.Context, body *extprocv3.HttpBody) (res *extprocv3.ProcessingResponse, err error) {
+func (e *embeddingsProcessorUpstreamFilter) ProcessResponseBody(ctx context.Context, body *extprocv3.HttpBody) (res *extprocv3.ProcessingResponse, err error) {
 	recordRequestCompletionErr := false
 	defer func() {
 		if err != nil || recordRequestCompletionErr {
@@ -385,13 +385,13 @@ func (e *embeddingsProcessorUpstreamFilter[T]) ProcessResponseBody(ctx context.C
 }
 
 // SetBackend implements [Processor.SetBackend].
-func (e *embeddingsProcessorUpstreamFilter[T]) SetBackend(ctx context.Context, b *filterapi.Backend, backendHandler filterapi.BackendAuthHandler, routeProcessor Processor) (err error) {
+func (e *embeddingsProcessorUpstreamFilter) SetBackend(ctx context.Context, b *filterapi.Backend, backendHandler filterapi.BackendAuthHandler, routeProcessor Processor) (err error) {
 	defer func() {
 		if err != nil {
 			e.metrics.RecordRequestCompletion(ctx, false, e.requestHeaders)
 		}
 	}()
-	rp, ok := routeProcessor.(*embeddingsProcessorRouterFilter[T])
+	rp, ok := routeProcessor.(*embeddingsProcessorRouterFilter)
 	if !ok {
 		panic("BUG: expected routeProcessor to be of type *embeddingsProcessorRouterFilter")
 	}
@@ -420,27 +420,26 @@ func (e *embeddingsProcessorUpstreamFilter[T]) SetBackend(ctx context.Context, b
 }
 
 // convertToEmbeddingCompletionRequest converts any EmbeddingRequest to EmbeddingCompletionRequest for compatibility
-func convertToEmbeddingCompletionRequest[T openai.EmbeddingRequest](req *T) *openai.EmbeddingCompletionRequest {
-	switch r := any(*req).(type) {
-	case openai.EmbeddingCompletionRequest:
-		return &r
-	case openai.EmbeddingChatRequest:
+func convertToEmbeddingCompletionRequest(req *openai.EmbeddingRequest) *openai.EmbeddingCompletionRequest {
+	if req.OfCompletion != nil {
+		return req.OfCompletion
+	} else if req.OfChat != nil {
 		// Convert EmbeddingChatRequest to EmbeddingCompletionRequest by flattening messages to input
 		// This is a simplified conversion - in practice you might need more sophisticated logic
 		return &openai.EmbeddingCompletionRequest{
-			Model:          r.Model,
+			Model:          req.OfChat.Model,
 			Input:          openai.EmbeddingRequestInput{Value: "converted_from_chat"}, // Simplified
-			EncodingFormat: r.EncodingFormat,
-			Dimensions:     r.Dimensions,
-			User:           r.User,
+			EncodingFormat: req.OfChat.EncodingFormat,
+			Dimensions:     req.OfChat.Dimensions,
+			User:           req.OfChat.User,
 		}
-	default:
+	} else {
 		return &openai.EmbeddingCompletionRequest{}
 	}
 }
 
-func parseOpenAIEmbeddingBody[T openai.EmbeddingRequest](body *extprocv3.HttpBody) (modelName string, rb *T, err error) {
-	var openAIReq T
+func parseOpenAIEmbeddingBody(body *extprocv3.HttpBody) (modelName string, rb *openai.EmbeddingRequest, err error) {
+	var openAIReq openai.EmbeddingRequest
 	if err := json.Unmarshal(body.Body, &openAIReq); err != nil {
 		return "", nil, fmt.Errorf("failed to unmarshal body: %w", err)
 	}

internal/extproc/embeddings_processor_test.go

@@ -371,7 +371,11 @@ func Test_embeddingsProcessorUpstreamFilter_ProcessRequestHeaders(t *testing.T)
 
 func TestEmbeddings_ProcessRequestHeaders_SetsRequestModel(t *testing.T) {
 	headers := map[string]string{":path": "/v1/embeddings", internalapi.ModelNameHeaderKeyDefault: "header-model"}
-	body := openai.EmbeddingRequest{Model: "body-model"}
+	body := openai.EmbeddingRequest{
+		OfCompletion: &openai.EmbeddingCompletionRequest{
+			Model: "body-model",
+		},
+	}
 	raw, _ := json.Marshal(body)
 	mm := &mockMetrics{}
 	p := &embeddingsProcessorUpstreamFilter{
@@ -395,8 +399,10 @@ func TestEmbeddings_ProcessResponseBody_OverridesHeaderModelWithResponseModel(t
 	const modelKey = internalapi.ModelNameHeaderKeyDefault
 	headers := map[string]string{":path": "/v1/embeddings", modelKey: "header-model"}
 	body := openai.EmbeddingRequest{
-		Model: "body-model",
-		Input: openai.EmbeddingRequestInput{Value: "test"},
+		OfCompletion: &openai.EmbeddingCompletionRequest{
+			Model: "body-model",
+			Input: openai.EmbeddingRequestInput{Value: "test"},
+		},
 	}
 	raw, _ := json.Marshal(body)
 	mm := &mockMetrics{}
@@ -454,8 +460,9 @@ func TestEmbeddings_ParseBody(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, "text-embedding-ada-002", modelName)
 		require.NotNil(t, rb)
-		require.Equal(t, "text-embedding-ada-002", rb.Model)
-		require.Equal(t, "test input", rb.Input.Value)
+		require.NotNil(t, rb.OfCompletion, "should be a completion request")
+		require.Equal(t, "text-embedding-ada-002", rb.OfCompletion.Model)
+		require.Equal(t, "test input", rb.OfCompletion.Input.Value)
 	})
 	t.Run("error", func(t *testing.T) {
 		modelName, rb, err := parseOpenAIEmbeddingBody(&extprocv3.HttpBody{})
@@ -680,7 +687,9 @@ func TestEmbeddingsProcessorUpstreamFilter_ProcessRequestHeaders_WithBodyMutatio
 		}
 
 		requestBody := &openai.EmbeddingRequest{
-			Model: "text-embedding-ada-002",
+			OfCompletion: &openai.EmbeddingCompletionRequest{
+				Model: "text-embedding-ada-002",
+			},
 		}
 		requestBodyRaw := []byte(`{"model": "text-embedding-ada-002", "input": "Hello world", "encoding_format": "float", "dimensions": 1536}`)
 
@@ -757,7 +766,9 @@ func TestEmbeddingsProcessorUpstreamFilter_ProcessRequestHeaders_WithBodyMutatio
 
 		originalRequestBodyRaw := []byte(`{"model": "text-embedding-ada-002", "input": "Original input", "encoding_format": "float"}`)
 		requestBody := &openai.EmbeddingRequest{
-			Model: "text-embedding-ada-002",
+			OfCompletion: &openai.EmbeddingCompletionRequest{
+				Model: "text-embedding-ada-002",
+			},
 		}
 
 		p := &embeddingsProcessorUpstreamFilter{