feat(vertex): support OpenAI-compatible embeddings

relay/channel/vertex/adaptor.go

@@ -50,6 +50,8 @@ var claudeModelMap = map[string]string{
 
 const anthropicVersion = "vertex-2023-10-16"
 
+const geminiEmbedding001MaxDimensions = 3072
+
 type Adaptor struct {
 	RequestMode        int
 	AccountCredentials Credentials
@@ -170,6 +172,10 @@ func (a *Adaptor) getRequestUrl(info *relaycommon.RelayInfo, modelName, suffix s
 func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
 	suffix := ""
 	if a.RequestMode == RequestModeGemini {
+		if info.RelayMode == constant.RelayModeEmbeddings {
+			return a.getRequestUrl(info, info.UpstreamModelName, "predict")
+		}
+
 		if model_setting.GetGeminiSettings().ThinkingAdapterEnabled &&
 			!model_setting.ShouldPreserveThinkingSuffix(info.OriginModelName) {
 			// 新增逻辑：处理 -thinking-<budget> 格式
@@ -323,8 +329,75 @@ func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dt
 }
 
 func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
-	//TODO implement me
-	return nil, errors.New("not implemented")
+	input, inputErr := parseSingleVertexEmbeddingInput(request.Input)
+	if inputErr != nil {
+		return nil, invalidVertexEmbeddingRequest(inputErr)
+	}
+	if request.EncodingFormat != "" && request.EncodingFormat != "float" {
+		return nil, invalidVertexEmbeddingRequest(
+			fmt.Errorf("Vertex embedding does not support encoding_format %q", request.EncodingFormat),
+		)
+	}
+	if request.Dimensions != nil && *request.Dimensions <= 0 {
+		return nil, invalidVertexEmbeddingRequest(errors.New("Vertex embedding dimensions must be greater than zero"))
+	}
+	if request.Dimensions != nil && info.UpstreamModelName == "gemini-embedding-001" && *request.Dimensions > geminiEmbedding001MaxDimensions {
+		return nil, invalidVertexEmbeddingRequest(
+			fmt.Errorf(
+				"Vertex model %s supports at most %d embedding dimensions",
+				info.UpstreamModelName,
+				geminiEmbedding001MaxDimensions,
+			),
+		)
+	}
+
+	vertexRequest := &VertexEmbeddingRequest{
+		Instances: []VertexEmbeddingInstance{{Content: input}},
+	}
+	if request.Dimensions != nil {
+		vertexRequest.Parameters = &VertexEmbeddingParameters{
+			OutputDimensionality: request.Dimensions,
+		}
+	}
+	return vertexRequest, nil
+}
+
+func parseSingleVertexEmbeddingInput(input any) (string, error) {
+	var value any
+	switch typedInput := input.(type) {
+	case string:
+		value = typedInput
+	case []any:
+		if len(typedInput) != 1 {
+			return "", fmt.Errorf("Vertex embedding requires exactly one input, got %d", len(typedInput))
+		}
+		value = typedInput[0]
+	case []string:
+		if len(typedInput) != 1 {
+			return "", fmt.Errorf("Vertex embedding requires exactly one input, got %d", len(typedInput))
+		}
+		value = typedInput[0]
+	default:
+		return "", fmt.Errorf("Vertex embedding input must be a string or a single-element string array, got %T", input)
+	}
+
+	text, ok := value.(string)
+	if !ok {
+		return "", fmt.Errorf("Vertex embedding input array must contain a string, got %T", value)
+	}
+	if strings.TrimSpace(text) == "" {
+		return "", errors.New("Vertex embedding input must not be empty")
+	}
+	return text, nil
+}
+
+func invalidVertexEmbeddingRequest(err error) *types.NewAPIError {
+	return types.NewErrorWithStatusCode(
+		err,
+		types.ErrorCodeInvalidRequest,
+		http.StatusBadRequest,
+		types.ErrOptionWithSkipRetry(),
+	)
 }
 
 func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
@@ -337,6 +410,10 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request
 }
 
 func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
+	if info.RelayMode == constant.RelayModeEmbeddings {
+		return VertexEmbeddingHandler(c, info, resp)
+	}
+
 	claudeAdaptor := claude.Adaptor{}
 	if info.IsStream {
 		switch a.RequestMode {

relay/channel/vertex/dto.go

@@ -23,6 +23,37 @@ type VertexAIClaudeRequest struct {
 	//Metadata         json.RawMessage     `json:"metadata,omitempty"`
 }
 
+type VertexEmbeddingInstance struct {
+	Content string `json:"content"`
+}
+
+type VertexEmbeddingParameters struct {
+	OutputDimensionality *int `json:"outputDimensionality,omitempty"`
+}
+
+type VertexEmbeddingRequest struct {
+	Instances  []VertexEmbeddingInstance  `json:"instances"`
+	Parameters *VertexEmbeddingParameters `json:"parameters,omitempty"`
+}
+
+type VertexEmbeddingStatistics struct {
+	TokenCount int  `json:"token_count"`
+	Truncated  bool `json:"truncated"`
+}
+
+type VertexEmbedding struct {
+	Values     []float64                 `json:"values"`
+	Statistics VertexEmbeddingStatistics `json:"statistics"`
+}
+
+type VertexEmbeddingPrediction struct {
+	Embeddings VertexEmbedding `json:"embeddings"`
+}
+
+type VertexEmbeddingResponse struct {
+	Predictions []VertexEmbeddingPrediction `json:"predictions"`
+}
+
 func copyRequest(req *dto.ClaudeRequest, version string) *VertexAIClaudeRequest {
 	return &VertexAIClaudeRequest{
 		AnthropicVersion: version,

relay/channel/vertex/embedding.go

@@ -0,0 +1,64 @@
+package vertex
+
+import (
+	"errors"
+	"io"
+	"net/http"
+
+	"github.com/QuantumNous/new-api/common"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/QuantumNous/new-api/relaykit/dto"
+	"github.com/QuantumNous/new-api/relaykit/types"
+	"github.com/QuantumNous/new-api/service"
+
+	"github.com/gin-gonic/gin"
+)
+
+func VertexEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
+	defer service.CloseResponseBodyGracefully(resp)
+
+	responseBody, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
+	}
+
+	var vertexResponse VertexEmbeddingResponse
+	if err := common.Unmarshal(responseBody, &vertexResponse); err != nil {
+		return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
+	}
+	if len(vertexResponse.Predictions) != 1 {
+		return nil, types.NewOpenAIError(errors.New("Vertex embedding response must contain exactly one prediction"), types.ErrorCodeBadResponseBody, http.StatusBadGateway)
+	}
+
+	prediction := vertexResponse.Predictions[0]
+	if len(prediction.Embeddings.Values) == 0 {
+		return nil, types.NewOpenAIError(errors.New("Vertex embedding response contains an empty vector"), types.ErrorCodeBadResponseBody, http.StatusBadGateway)
+	}
+
+	promptTokens := prediction.Embeddings.Statistics.TokenCount
+	if promptTokens <= 0 {
+		promptTokens = info.GetEstimatePromptTokens()
+	}
+	usage := dto.Usage{
+		PromptTokens: promptTokens,
+		TotalTokens:  promptTokens,
+		InputTokens:  promptTokens,
+	}
+	openAIResponse := dto.OpenAIEmbeddingResponse{
+		Object: "list",
+		Data: []dto.OpenAIEmbeddingResponseItem{{
+			Object:    "embedding",
+			Index:     0,
+			Embedding: prediction.Embeddings.Values,
+		}},
+		Model: info.UpstreamModelName,
+		Usage: usage,
+	}
+
+	jsonResponse, err := common.Marshal(openAIResponse)
+	if err != nil {
+		return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
+	}
+	service.IOCopyBytesGracefully(c, resp, jsonResponse)
+	return &usage, nil
+}