Merge branch 'main' into pr/Laisky/25

relay/controller/audio.go

@@ -0,0 +1,333 @@
+package controller
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"mime/multipart"
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+	"github.com/pkg/errors"
+	"github.com/songquanpeng/one-api/common"
+	"github.com/songquanpeng/one-api/common/client"
+	"github.com/songquanpeng/one-api/common/ctxkey"
+	"github.com/songquanpeng/one-api/common/helper"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/model"
+	"github.com/songquanpeng/one-api/relay/adaptor/openai"
+	"github.com/songquanpeng/one-api/relay/billing"
+	"github.com/songquanpeng/one-api/relay/billing/ratio"
+	billingratio "github.com/songquanpeng/one-api/relay/billing/ratio"
+	"github.com/songquanpeng/one-api/relay/channeltype"
+	"github.com/songquanpeng/one-api/relay/meta"
+	relaymodel "github.com/songquanpeng/one-api/relay/model"
+	"github.com/songquanpeng/one-api/relay/relaymode"
+)
+
+type commonAudioRequest struct {
+	File *multipart.FileHeader `form:"file" binding:"required"`
+}
+
+func countAudioTokens(c *gin.Context) (int, error) {
+	body, err := common.GetRequestBody(c)
+	if err != nil {
+		return 0, errors.WithStack(err)
+	}
+
+	reqBody := new(commonAudioRequest)
+	c.Request.Body = io.NopCloser(bytes.NewReader(body))
+	if err = c.ShouldBind(reqBody); err != nil {
+		return 0, errors.WithStack(err)
+	}
+
+	reqFp, err := reqBody.File.Open()
+	if err != nil {
+		return 0, errors.WithStack(err)
+	}
+	defer reqFp.Close()
+
+	ctxMeta := meta.GetByContext(c)
+
+	return helper.GetAudioTokens(c.Request.Context(),
+		reqFp,
+		ratio.GetAudioPromptTokensPerSecond(ctxMeta.ActualModelName))
+}
+
+func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatusCode {
+	ctx := c.Request.Context()
+	meta := meta.GetByContext(c)
+	audioModel := "whisper-1"
+
+	tokenId := c.GetInt(ctxkey.TokenId)
+	channelType := c.GetInt(ctxkey.Channel)
+	channelId := c.GetInt(ctxkey.ChannelId)
+	userId := c.GetInt(ctxkey.Id)
+	// group := c.GetString(ctxkey.Group)
+	tokenName := c.GetString(ctxkey.TokenName)
+
+	var ttsRequest openai.TextToSpeechRequest
+	if relayMode == relaymode.AudioSpeech {
+		// Read JSON
+		err := common.UnmarshalBodyReusable(c, &ttsRequest)
+		// Check if JSON is valid
+		if err != nil {
+			return openai.ErrorWrapper(err, "invalid_json", http.StatusBadRequest)
+		}
+		audioModel = ttsRequest.Model
+		// Check if text is too long 4096
+		if len(ttsRequest.Input) > 4096 {
+			return openai.ErrorWrapper(errors.New("input is too long (over 4096 characters)"), "text_too_long", http.StatusBadRequest)
+		}
+	}
+
+	modelRatio := billingratio.GetModelRatio(audioModel, channelType)
+	// groupRatio := billingratio.GetGroupRatio(group)
+	groupRatio := c.GetFloat64(ctxkey.ChannelRatio)
+	ratio := modelRatio * groupRatio
+	var quota int64
+	var preConsumedQuota int64
+	switch relayMode {
+	case relaymode.AudioSpeech:
+		preConsumedQuota = int64(float64(len(ttsRequest.Input)) * ratio)
+		quota = preConsumedQuota
+	case relaymode.AudioTranscription,
+		relaymode.AudioTranslation:
+		audioTokens, err := countAudioTokens(c)
+		if err != nil {
+			return openai.ErrorWrapper(err, "count_audio_tokens_failed", http.StatusInternalServerError)
+		}
+
+		preConsumedQuota = int64(float64(audioTokens) * ratio)
+		quota = preConsumedQuota
+	default:
+		return openai.ErrorWrapper(errors.New("unexpected_relay_mode"), "unexpected_relay_mode", http.StatusInternalServerError)
+	}
+
+	tokenQuota := c.GetInt64(ctxkey.TokenQuota)
+	tokenQuotaUnlimited := c.GetBool(ctxkey.TokenQuotaUnlimited)
+	userQuota, err := model.CacheGetUserQuota(ctx, userId)
+	if err != nil {
+		return openai.ErrorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError)
+	}
+
+	// Check if user quota is enough
+	if userQuota-preConsumedQuota < 0 {
+		return openai.ErrorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden)
+	}
+	err = model.CacheDecreaseUserQuota(userId, preConsumedQuota)
+	if err != nil {
+		return openai.ErrorWrapper(err, "decrease_user_quota_failed", http.StatusInternalServerError)
+	}
+	if userQuota > 100*preConsumedQuota &&
+		(tokenQuotaUnlimited || tokenQuota > 100*preConsumedQuota) {
+		// in this case, we do not pre-consume quota
+		// because the user has enough quota
+		preConsumedQuota = 0
+	}
+	if preConsumedQuota > 0 {
+		err := model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
+		if err != nil {
+			return openai.ErrorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden)
+		}
+	}
+	succeed := false
+	defer func() {
+		if succeed {
+			return
+		}
+		if preConsumedQuota > 0 {
+			// we need to roll back the pre-consumed quota
+			defer func(ctx context.Context) {
+				go func() {
+					// negative means add quota back for token & user
+					err := model.PostConsumeTokenQuota(tokenId, -preConsumedQuota)
+					if err != nil {
+						logger.Error(ctx, fmt.Sprintf("error rollback pre-consumed quota: %s", err.Error()))
+					}
+				}()
+			}(c.Request.Context())
+		}
+	}()
+
+	// map model name
+	modelMapping := c.GetStringMapString(ctxkey.ModelMapping)
+	if modelMapping != nil && modelMapping[audioModel] != "" {
+		audioModel = modelMapping[audioModel]
+	}
+
+	baseURL := channeltype.ChannelBaseURLs[channelType]
+	requestURL := c.Request.URL.String()
+	if c.GetString(ctxkey.BaseURL) != "" {
+		baseURL = c.GetString(ctxkey.BaseURL)
+	}
+
+	fullRequestURL := openai.GetFullRequestURL(baseURL, requestURL, channelType)
+	if channelType == channeltype.Azure {
+		apiVersion := meta.Config.APIVersion
+		if relayMode == relaymode.AudioTranscription {
+			// https://learn.microsoft.com/en-us/azure/ai-services/openai/whisper-quickstart?tabs=command-line#rest-api
+			fullRequestURL = fmt.Sprintf("%s/openai/deployments/%s/audio/transcriptions?api-version=%s", baseURL, audioModel, apiVersion)
+		} else if relayMode == relaymode.AudioSpeech {
+			// https://learn.microsoft.com/en-us/azure/ai-services/openai/text-to-speech-quickstart?tabs=command-line#rest-api
+			fullRequestURL = fmt.Sprintf("%s/openai/deployments/%s/audio/speech?api-version=%s", baseURL, audioModel, apiVersion)
+		}
+	}
+
+	requestBody := &bytes.Buffer{}
+	_, err = io.Copy(requestBody, c.Request.Body)
+	if err != nil {
+		return openai.ErrorWrapper(err, "new_request_body_failed", http.StatusInternalServerError)
+	}
+	c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody.Bytes()))
+	// responseFormat := c.DefaultPostForm("response_format", "json")
+
+	req, err := http.NewRequest(c.Request.Method, fullRequestURL, requestBody)
+	if err != nil {
+		return openai.ErrorWrapper(err, "new_request_failed", http.StatusInternalServerError)
+	}
+
+	if (relayMode == relaymode.AudioTranscription || relayMode == relaymode.AudioSpeech) && channelType == channeltype.Azure {
+		// https://learn.microsoft.com/en-us/azure/ai-services/openai/whisper-quickstart?tabs=command-line#rest-api
+		apiKey := c.Request.Header.Get("Authorization")
+		apiKey = strings.TrimPrefix(apiKey, "Bearer ")
+		req.Header.Set("api-key", apiKey)
+		req.ContentLength = c.Request.ContentLength
+	} else {
+		req.Header.Set("Authorization", c.Request.Header.Get("Authorization"))
+	}
+	req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
+	req.Header.Set("Accept", c.Request.Header.Get("Accept"))
+
+	resp, err := client.HTTPClient.Do(req)
+	if err != nil {
+		return openai.ErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
+	}
+
+	err = req.Body.Close()
+	if err != nil {
+		return openai.ErrorWrapper(err, "close_request_body_failed", http.StatusInternalServerError)
+	}
+	err = c.Request.Body.Close()
+	if err != nil {
+		return openai.ErrorWrapper(err, "close_request_body_failed", http.StatusInternalServerError)
+	}
+
+	// https://github.com/Laisky/one-api/pull/21
+	// Commenting out the following code because Whisper's transcription
+	// only charges for the length of the input audio, not for the output.
+	// -------------------------------------
+	// if relayMode != relaymode.AudioSpeech {
+	// 	responseBody, err := io.ReadAll(resp.Body)
+	// 	if err != nil {
+	// 		return openai.ErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
+	// 	}
+	// 	err = resp.Body.Close()
+	// 	if err != nil {
+	// 		return openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError)
+	// 	}
+
+	// 	var openAIErr openai.SlimTextResponse
+	// 	if err = json.Unmarshal(responseBody, &openAIErr); err == nil {
+	// 		if openAIErr.Error.Message != "" {
+	// 			return openai.ErrorWrapper(errors.Errorf("type %s, code %v, message %s", openAIErr.Error.Type, openAIErr.Error.Code, openAIErr.Error.Message), "request_error", http.StatusInternalServerError)
+	// 		}
+	// 	}
+
+	// 	var text string
+	// 	switch responseFormat {
+	// 	case "json":
+	// 		text, err = getTextFromJSON(responseBody)
+	// 	case "text":
+	// 		text, err = getTextFromText(responseBody)
+	// 	case "srt":
+	// 		text, err = getTextFromSRT(responseBody)
+	// 	case "verbose_json":
+	// 		text, err = getTextFromVerboseJSON(responseBody)
+	// 	case "vtt":
+	// 		text, err = getTextFromVTT(responseBody)
+	// 	default:
+	// 		return openai.ErrorWrapper(errors.New("unexpected_response_format"), "unexpected_response_format", http.StatusInternalServerError)
+	// 	}
+	// 	if err != nil {
+	// 		return openai.ErrorWrapper(err, "get_text_from_body_err", http.StatusInternalServerError)
+	// 	}
+	// 	quota = int64(openai.CountTokenText(text, audioModel))
+	// 	resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
+	// }
+
+	if resp.StatusCode != http.StatusOK {
+		return RelayErrorHandler(resp)
+	}
+
+	succeed = true
+	quotaDelta := quota - preConsumedQuota
+	defer func(ctx context.Context) {
+		go billing.PostConsumeQuota(ctx, tokenId, quotaDelta, quota, userId, channelId, modelRatio, groupRatio, audioModel, tokenName)
+	}(c.Request.Context())
+
+	for k, v := range resp.Header {
+		c.Writer.Header().Set(k, v[0])
+	}
+	c.Writer.WriteHeader(resp.StatusCode)
+
+	_, err = io.Copy(c.Writer, resp.Body)
+	if err != nil {
+		return openai.ErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError)
+	}
+	err = resp.Body.Close()
+	if err != nil {
+		return openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError)
+	}
+	return nil
+}
+
+func getTextFromVTT(body []byte) (string, error) {
+	return getTextFromSRT(body)
+}
+
+func getTextFromVerboseJSON(body []byte) (string, error) {
+	var whisperResponse openai.WhisperVerboseJSONResponse
+	if err := json.Unmarshal(body, &whisperResponse); err != nil {
+		return "", errors.Wrap(err, "unmarshal_response_body_failed")
+	}
+
+	return whisperResponse.Text, nil
+}
+
+func getTextFromSRT(body []byte) (string, error) {
+	scanner := bufio.NewScanner(strings.NewReader(string(body)))
+	var builder strings.Builder
+	var textLine bool
+	for scanner.Scan() {
+		line := scanner.Text()
+		if textLine {
+			builder.WriteString(line)
+			textLine = false
+			continue
+		} else if strings.Contains(line, "-->") {
+			textLine = true
+			continue
+		}
+	}
+	if err := scanner.Err(); err != nil {
+		return "", err
+	}
+	return builder.String(), nil
+}
+
+func getTextFromText(body []byte) (string, error) {
+	return strings.TrimSuffix(string(body), "\n"), nil
+}
+
+func getTextFromJSON(body []byte) (string, error) {
+	var whisperResponse openai.WhisperJSONResponse
+	if err := json.Unmarshal(body, &whisperResponse); err != nil {
+		return "", errors.Wrap(err, "unmarshal_response_body_failed")
+	}
+	return whisperResponse.Text, nil
+}

relay/controller/error.go

@@ -0,0 +1,101 @@
+package controller
+
+import (
+	"encoding/json"
+	"fmt"
+	"github.com/songquanpeng/one-api/common/config"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/relay/model"
+	"io"
+	"net/http"
+	"strconv"
+)
+
+type GeneralErrorResponse struct {
+	Error    model.Error `json:"error"`
+	Message  string      `json:"message"`
+	Msg      string      `json:"msg"`
+	Err      string      `json:"err"`
+	ErrorMsg string      `json:"error_msg"`
+	Header   struct {
+		Message string `json:"message"`
+	} `json:"header"`
+	Response struct {
+		Error struct {
+			Message string `json:"message"`
+		} `json:"error"`
+	} `json:"response"`
+}
+
+func (e GeneralErrorResponse) ToMessage() string {
+	if e.Error.Message != "" {
+		return e.Error.Message
+	}
+	if e.Message != "" {
+		return e.Message
+	}
+	if e.Msg != "" {
+		return e.Msg
+	}
+	if e.Err != "" {
+		return e.Err
+	}
+	if e.ErrorMsg != "" {
+		return e.ErrorMsg
+	}
+	if e.Header.Message != "" {
+		return e.Header.Message
+	}
+	if e.Response.Error.Message != "" {
+		return e.Response.Error.Message
+	}
+	return ""
+}
+
+func RelayErrorHandler(resp *http.Response) (ErrorWithStatusCode *model.ErrorWithStatusCode) {
+	if resp == nil {
+		return &model.ErrorWithStatusCode{
+			StatusCode: 500,
+			Error: model.Error{
+				Message: "resp is nil",
+				Type:    "upstream_error",
+				Code:    "bad_response",
+			},
+		}
+	}
+	ErrorWithStatusCode = &model.ErrorWithStatusCode{
+		StatusCode: resp.StatusCode,
+		Error: model.Error{
+			Message: "",
+			Type:    "upstream_error",
+			Code:    "bad_response_status_code",
+			Param:   strconv.Itoa(resp.StatusCode),
+		},
+	}
+	responseBody, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return
+	}
+	if config.DebugEnabled {
+		logger.SysLog(fmt.Sprintf("error happened, status code: %d, response: \n%s", resp.StatusCode, string(responseBody)))
+	}
+	err = resp.Body.Close()
+	if err != nil {
+		return
+	}
+	var errResponse GeneralErrorResponse
+	err = json.Unmarshal(responseBody, &errResponse)
+	if err != nil {
+		return
+	}
+	if errResponse.Error.Message != "" {
+		// OpenAI format error, so we override the default one
+		ErrorWithStatusCode.Error = errResponse.Error
+	} else {
+		ErrorWithStatusCode.Error.Message = errResponse.ToMessage()
+	}
+	if ErrorWithStatusCode.Error.Message == "" {
+		ErrorWithStatusCode.Error.Message = fmt.Sprintf("bad response status code %d", resp.StatusCode)
+	}
+	return
+}

relay/controller/helper.go

@@ -0,0 +1,193 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"math"
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+	"github.com/pkg/errors"
+	"github.com/songquanpeng/one-api/common"
+	"github.com/songquanpeng/one-api/common/config"
+	"github.com/songquanpeng/one-api/common/ctxkey"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/model"
+	"github.com/songquanpeng/one-api/relay/adaptor/openai"
+	billingratio "github.com/songquanpeng/one-api/relay/billing/ratio"
+	"github.com/songquanpeng/one-api/relay/channeltype"
+	"github.com/songquanpeng/one-api/relay/constant/role"
+	"github.com/songquanpeng/one-api/relay/controller/validator"
+	"github.com/songquanpeng/one-api/relay/meta"
+	relaymodel "github.com/songquanpeng/one-api/relay/model"
+	"github.com/songquanpeng/one-api/relay/relaymode"
+)
+
+func getAndValidateTextRequest(c *gin.Context, relayMode int) (*relaymodel.GeneralOpenAIRequest, error) {
+	textRequest := &relaymodel.GeneralOpenAIRequest{}
+	err := common.UnmarshalBodyReusable(c, textRequest)
+	if err != nil {
+		return nil, err
+	}
+	if relayMode == relaymode.Moderations && textRequest.Model == "" {
+		textRequest.Model = "text-moderation-latest"
+	}
+	if relayMode == relaymode.Embeddings && textRequest.Model == "" {
+		textRequest.Model = c.Param("model")
+	}
+	err = validator.ValidateTextRequest(textRequest, relayMode)
+	if err != nil {
+		return nil, err
+	}
+	return textRequest, nil
+}
+
+func getPromptTokens(ctx context.Context, textRequest *relaymodel.GeneralOpenAIRequest, relayMode int) int {
+	switch relayMode {
+	case relaymode.ChatCompletions:
+		return openai.CountTokenMessages(ctx, textRequest.Messages, textRequest.Model)
+	case relaymode.Completions:
+		return openai.CountTokenInput(textRequest.Prompt, textRequest.Model)
+	case relaymode.Moderations:
+		return openai.CountTokenInput(textRequest.Input, textRequest.Model)
+	}
+	return 0
+}
+
+func getPreConsumedQuota(textRequest *relaymodel.GeneralOpenAIRequest, promptTokens int, ratio float64) int64 {
+	preConsumedTokens := config.PreConsumedQuota + int64(promptTokens)
+	if textRequest.MaxTokens != 0 {
+		preConsumedTokens += int64(textRequest.MaxTokens)
+	}
+	return int64(float64(preConsumedTokens) * ratio)
+}
+
+func preConsumeQuota(c *gin.Context, textRequest *relaymodel.GeneralOpenAIRequest, promptTokens int, ratio float64, meta *meta.Meta) (int64, *relaymodel.ErrorWithStatusCode) {
+	preConsumedQuota := getPreConsumedQuota(textRequest, promptTokens, ratio)
+
+	tokenQuota := c.GetInt64(ctxkey.TokenQuota)
+	tokenQuotaUnlimited := c.GetBool(ctxkey.TokenQuotaUnlimited)
+	userQuota, err := model.CacheGetUserQuota(c.Request.Context(), meta.UserId)
+	if err != nil {
+		return preConsumedQuota, openai.ErrorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError)
+	}
+	if userQuota-preConsumedQuota < 0 {
+		return preConsumedQuota, openai.ErrorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden)
+	}
+	err = model.CacheDecreaseUserQuota(meta.UserId, preConsumedQuota)
+	if err != nil {
+		return preConsumedQuota, openai.ErrorWrapper(err, "decrease_user_quota_failed", http.StatusInternalServerError)
+	}
+	if userQuota > 100*preConsumedQuota &&
+		(tokenQuotaUnlimited || tokenQuota > 100*preConsumedQuota) {
+		// in this case, we do not pre-consume quota
+		// because the user and token have enough quota
+		preConsumedQuota = 0
+		logger.Info(c.Request.Context(), fmt.Sprintf("user %d has enough quota %d, trusted and no need to pre-consume", meta.UserId, userQuota))
+	}
+	if preConsumedQuota > 0 {
+		err := model.PreConsumeTokenQuota(meta.TokenId, preConsumedQuota)
+		if err != nil {
+			return preConsumedQuota, openai.ErrorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden)
+		}
+	}
+	return preConsumedQuota, nil
+}
+
+func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *meta.Meta, textRequest *relaymodel.GeneralOpenAIRequest, ratio float64, preConsumedQuota int64, modelRatio float64, groupRatio float64, systemPromptReset bool) (quota int64) {
+	if usage == nil {
+		logger.Error(ctx, "usage is nil, which is unexpected")
+		return
+	}
+
+	completionRatio := billingratio.GetCompletionRatio(textRequest.Model, meta.ChannelType)
+	promptTokens := usage.PromptTokens
+	completionTokens := usage.CompletionTokens
+	quota = int64(math.Ceil((float64(promptTokens) + float64(completionTokens)*completionRatio) * ratio))
+	if ratio != 0 && quota <= 0 {
+		quota = 1
+	}
+	totalTokens := promptTokens + completionTokens
+	if totalTokens == 0 {
+		// in this case, must be some error happened
+		// we cannot just return, because we may have to return the pre-consumed quota
+		quota = 0
+	}
+	quotaDelta := quota - preConsumedQuota
+	err := model.PostConsumeTokenQuota(meta.TokenId, quotaDelta)
+	if err != nil {
+		logger.Error(ctx, "error consuming token remain quota: "+err.Error())
+	}
+	err = model.CacheUpdateUserQuota(ctx, meta.UserId)
+	if err != nil {
+		logger.Error(ctx, "error update user quota cache: "+err.Error())
+	}
+	var extraLog string
+	if systemPromptReset {
+		extraLog = " （注意系统提示词已被重置）"
+	}
+	logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f，补全倍率 %.2f%s", modelRatio, groupRatio, completionRatio, extraLog)
+	model.RecordConsumeLog(ctx, meta.UserId, meta.ChannelId, promptTokens, completionTokens, textRequest.Model, meta.TokenName, quota, logContent)
+	model.UpdateUserUsedQuotaAndRequestCount(meta.UserId, quota)
+	model.UpdateChannelUsedQuota(meta.ChannelId, quota)
+
+	return quota
+}
+
+func getMappedModelName(modelName string, mapping map[string]string) (string, bool) {
+	if mapping == nil {
+		return modelName, false
+	}
+	mappedModelName := mapping[modelName]
+	if mappedModelName != "" {
+		return mappedModelName, true
+	}
+	return modelName, false
+}
+
+func isErrorHappened(meta *meta.Meta, resp *http.Response) bool {
+	if resp == nil {
+		if meta.ChannelType == channeltype.AwsClaude {
+			return false
+		}
+		return true
+	}
+	if resp.StatusCode != http.StatusOK &&
+		// replicate return 201 to create a task
+		resp.StatusCode != http.StatusCreated {
+		return true
+	}
+	if meta.ChannelType == channeltype.DeepL {
+		// skip stream check for deepl
+		return false
+	}
+
+	if meta.IsStream && strings.HasPrefix(resp.Header.Get("Content-Type"), "application/json") &&
+		// Even if stream mode is enabled, replicate will first return a task info in JSON format,
+		// requiring the client to request the stream endpoint in the task info
+		meta.ChannelType != channeltype.Replicate {
+		return true
+	}
+	return false
+}
+
+func setSystemPrompt(ctx context.Context, request *relaymodel.GeneralOpenAIRequest, prompt string) (reset bool) {
+	if prompt == "" {
+		return false
+	}
+	if len(request.Messages) == 0 {
+		return false
+	}
+	if request.Messages[0].Role == role.System {
+		request.Messages[0].Content = prompt
+		logger.Infof(ctx, "rewrite system prompt")
+		return true
+	}
+	request.Messages = append([]relaymodel.Message{{
+		Role:    role.System,
+		Content: prompt,
+	}}, request.Messages...)
+	logger.Infof(ctx, "add system prompt")
+	return true
+}

relay/controller/image.go

@@ -0,0 +1,254 @@
+package controller
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+	"github.com/pkg/errors"
+	"github.com/songquanpeng/one-api/common"
+	"github.com/songquanpeng/one-api/common/ctxkey"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/model"
+	"github.com/songquanpeng/one-api/relay"
+	"github.com/songquanpeng/one-api/relay/adaptor/openai"
+	"github.com/songquanpeng/one-api/relay/adaptor/replicate"
+	billingratio "github.com/songquanpeng/one-api/relay/billing/ratio"
+	"github.com/songquanpeng/one-api/relay/channeltype"
+	metalib "github.com/songquanpeng/one-api/relay/meta"
+	relaymodel "github.com/songquanpeng/one-api/relay/model"
+)
+
+func getImageRequest(c *gin.Context, _ int) (*relaymodel.ImageRequest, error) {
+	imageRequest := &relaymodel.ImageRequest{}
+	err := common.UnmarshalBodyReusable(c, imageRequest)
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	if imageRequest.N == 0 {
+		imageRequest.N = 1
+	}
+	if imageRequest.Size == "" {
+		imageRequest.Size = "1024x1024"
+	}
+	if imageRequest.Model == "" {
+		imageRequest.Model = "dall-e-2"
+	}
+	return imageRequest, nil
+}
+
+func isValidImageSize(model string, size string) bool {
+	if model == "cogview-3" || billingratio.ImageSizeRatios[model] == nil {
+		return true
+	}
+	_, ok := billingratio.ImageSizeRatios[model][size]
+	return ok
+}
+
+func isValidImagePromptLength(model string, promptLength int) bool {
+	maxPromptLength, ok := billingratio.ImagePromptLengthLimitations[model]
+	return !ok || promptLength <= maxPromptLength
+}
+
+func isWithinRange(element string, value int) bool {
+	amounts, ok := billingratio.ImageGenerationAmounts[element]
+	return !ok || (value >= amounts[0] && value <= amounts[1])
+}
+
+func getImageSizeRatio(model string, size string) float64 {
+	if ratio, ok := billingratio.ImageSizeRatios[model][size]; ok {
+		return ratio
+	}
+	return 1
+}
+
+func validateImageRequest(imageRequest *relaymodel.ImageRequest, _ *metalib.Meta) *relaymodel.ErrorWithStatusCode {
+	// check prompt length
+	if imageRequest.Prompt == "" {
+		return openai.ErrorWrapper(errors.New("prompt is required"), "prompt_missing", http.StatusBadRequest)
+	}
+
+	// model validation
+	if !isValidImageSize(imageRequest.Model, imageRequest.Size) {
+		return openai.ErrorWrapper(errors.New("size not supported for this image model"), "size_not_supported", http.StatusBadRequest)
+	}
+
+	if !isValidImagePromptLength(imageRequest.Model, len(imageRequest.Prompt)) {
+		return openai.ErrorWrapper(errors.New("prompt is too long"), "prompt_too_long", http.StatusBadRequest)
+	}
+
+	// Number of generated images validation
+	if !isWithinRange(imageRequest.Model, imageRequest.N) {
+		return openai.ErrorWrapper(errors.New("invalid value of n"), "n_not_within_range", http.StatusBadRequest)
+	}
+	return nil
+}
+
+func getImageCostRatio(imageRequest *relaymodel.ImageRequest) (float64, error) {
+	if imageRequest == nil {
+		return 0, errors.New("imageRequest is nil")
+	}
+	imageCostRatio := getImageSizeRatio(imageRequest.Model, imageRequest.Size)
+	if imageRequest.Quality == "hd" && imageRequest.Model == "dall-e-3" {
+		if imageRequest.Size == "1024x1024" {
+			imageCostRatio *= 2
+		} else {
+			imageCostRatio *= 1.5
+		}
+	}
+	return imageCostRatio, nil
+}
+
+func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatusCode {
+	ctx := c.Request.Context()
+	meta := metalib.GetByContext(c)
+	imageRequest, err := getImageRequest(c, meta.Mode)
+	if err != nil {
+		logger.Errorf(ctx, "getImageRequest failed: %s", err.Error())
+		return openai.ErrorWrapper(err, "invalid_image_request", http.StatusBadRequest)
+	}
+
+	// map model name
+	var isModelMapped bool
+	meta.OriginModelName = imageRequest.Model
+	imageRequest.Model, isModelMapped = getMappedModelName(imageRequest.Model, meta.ModelMapping)
+	meta.ActualModelName = imageRequest.Model
+	metalib.Set2Context(c, meta)
+
+	// model validation
+	bizErr := validateImageRequest(imageRequest, meta)
+	if bizErr != nil {
+		return bizErr
+	}
+
+	imageCostRatio, err := getImageCostRatio(imageRequest)
+	if err != nil {
+		return openai.ErrorWrapper(err, "get_image_cost_ratio_failed", http.StatusInternalServerError)
+	}
+
+	imageModel := imageRequest.Model
+	// Convert the original image model
+	imageRequest.Model, _ = getMappedModelName(imageRequest.Model, billingratio.ImageOriginModelName)
+	c.Set("response_format", imageRequest.ResponseFormat)
+
+	var requestBody io.Reader
+	if strings.ToLower(c.GetString(ctxkey.ContentType)) == "application/json" &&
+		isModelMapped || meta.ChannelType == channeltype.Azure { // make Azure channel request body
+		jsonStr, err := json.Marshal(imageRequest)
+		if err != nil {
+			return openai.ErrorWrapper(err, "marshal_image_request_failed", http.StatusInternalServerError)
+		}
+		requestBody = bytes.NewBuffer(jsonStr)
+	} else {
+		requestBody = c.Request.Body
+	}
+
+	adaptor := relay.GetAdaptor(meta.APIType)
+	if adaptor == nil {
+		return openai.ErrorWrapper(fmt.Errorf("invalid api type: %d", meta.APIType), "invalid_api_type", http.StatusBadRequest)
+	}
+	adaptor.Init(meta)
+
+	// these adaptors need to convert the request
+	switch meta.ChannelType {
+	case channeltype.Zhipu,
+		channeltype.Ali,
+		channeltype.VertextAI,
+		channeltype.Baidu:
+		finalRequest, err := adaptor.ConvertImageRequest(c, imageRequest)
+		if err != nil {
+			return openai.ErrorWrapper(err, "convert_image_request_failed", http.StatusInternalServerError)
+		}
+		jsonStr, err := json.Marshal(finalRequest)
+		if err != nil {
+			return openai.ErrorWrapper(err, "marshal_image_request_failed", http.StatusInternalServerError)
+		}
+		requestBody = bytes.NewBuffer(jsonStr)
+	case channeltype.Replicate:
+		finalRequest, err := replicate.ConvertImageRequest(c, imageRequest)
+		if err != nil {
+			return openai.ErrorWrapper(err, "convert_image_request_failed", http.StatusInternalServerError)
+		}
+		jsonStr, err := json.Marshal(finalRequest)
+		if err != nil {
+			return openai.ErrorWrapper(err, "marshal_image_request_failed", http.StatusInternalServerError)
+		}
+		requestBody = bytes.NewBuffer(jsonStr)
+	}
+
+	modelRatio := billingratio.GetModelRatio(imageModel, meta.ChannelType)
+	// groupRatio := billingratio.GetGroupRatio(meta.Group)
+	groupRatio := c.GetFloat64(ctxkey.ChannelRatio)
+
+	ratio := modelRatio * groupRatio
+	userQuota, err := model.CacheGetUserQuota(ctx, meta.UserId)
+
+	var quota int64
+	switch meta.ChannelType {
+	case channeltype.Replicate:
+		// replicate always return 1 image
+		quota = int64(ratio * imageCostRatio * 1000)
+	default:
+		quota = int64(ratio*imageCostRatio*1000) * int64(imageRequest.N)
+	}
+
+	if userQuota-quota < 0 {
+		return openai.ErrorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden)
+	}
+
+	// do request
+	resp, err := adaptor.DoRequest(c, meta, requestBody)
+	if err != nil {
+		logger.Errorf(ctx, "DoRequest failed: %s", err.Error())
+		return openai.ErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
+	}
+
+	defer func(ctx context.Context) {
+		if resp != nil &&
+			resp.StatusCode != http.StatusCreated && // replicate returns 201
+			resp.StatusCode != http.StatusOK {
+			return
+		}
+
+		err := model.PostConsumeTokenQuota(meta.TokenId, quota)
+		if err != nil {
+			logger.SysError("error consuming token remain quota: " + err.Error())
+		}
+		err = model.CacheUpdateUserQuota(ctx, meta.UserId)
+		if err != nil {
+			logger.SysError("error update user quota cache: " + err.Error())
+		}
+		if quota >= 0 {
+			tokenName := c.GetString(ctxkey.TokenName)
+			logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio)
+			model.RecordConsumeLog(ctx, meta.UserId, meta.ChannelId, 0, 0, imageRequest.Model, tokenName, quota, logContent)
+			model.UpdateUserUsedQuotaAndRequestCount(meta.UserId, quota)
+			channelId := c.GetInt(ctxkey.ChannelId)
+			model.UpdateChannelUsedQuota(channelId, quota)
+
+			// also update user request cost
+			docu := model.NewUserRequestCost(
+				c.GetInt(ctxkey.Id),
+				c.GetString(ctxkey.RequestId),
+				quota,
+			)
+			if err = docu.Insert(); err != nil {
+				logger.Errorf(c, "insert user request cost failed: %+v", err)
+			}
+		}
+	}(c.Request.Context())
+
+	// do response
+	_, respErr := adaptor.DoResponse(c, resp, meta)
+	if respErr != nil {
+		logger.Errorf(ctx, "respErr is not nil: %+v", respErr)
+		return respErr
+	}
+
+	return nil
+}

relay/controller/proxy.go

@@ -0,0 +1,41 @@
+// Package controller is a package for handling the relay controller
+package controller
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/relay"
+	"github.com/songquanpeng/one-api/relay/adaptor/openai"
+	"github.com/songquanpeng/one-api/relay/meta"
+	relaymodel "github.com/songquanpeng/one-api/relay/model"
+)
+
+// RelayProxyHelper is a helper function to proxy the request to the upstream service
+func RelayProxyHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatusCode {
+	ctx := c.Request.Context()
+	meta := meta.GetByContext(c)
+
+	adaptor := relay.GetAdaptor(meta.APIType)
+	if adaptor == nil {
+		return openai.ErrorWrapper(fmt.Errorf("invalid api type: %d", meta.APIType), "invalid_api_type", http.StatusBadRequest)
+	}
+	adaptor.Init(meta)
+
+	resp, err := adaptor.DoRequest(c, meta, c.Request.Body)
+	if err != nil {
+		logger.Errorf(ctx, "DoRequest failed: %s", err.Error())
+		return openai.ErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
+	}
+
+	// do response
+	_, respErr := adaptor.DoResponse(c, resp, meta)
+	if respErr != nil {
+		logger.Errorf(ctx, "respErr is not nil: %+v", respErr)
+		return respErr
+	}
+
+	return nil
+}

relay/controller/text.go

@@ -0,0 +1,144 @@
+package controller
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/pkg/errors"
+	"github.com/songquanpeng/one-api/common/config"
+	"github.com/songquanpeng/one-api/common/ctxkey"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/model"
+	"github.com/songquanpeng/one-api/relay"
+	"github.com/songquanpeng/one-api/relay/adaptor"
+	"github.com/songquanpeng/one-api/relay/adaptor/openai"
+	"github.com/songquanpeng/one-api/relay/apitype"
+	"github.com/songquanpeng/one-api/relay/billing"
+	billingratio "github.com/songquanpeng/one-api/relay/billing/ratio"
+	"github.com/songquanpeng/one-api/relay/channeltype"
+	"github.com/songquanpeng/one-api/relay/meta"
+	relaymodel "github.com/songquanpeng/one-api/relay/model"
+)
+
+func RelayTextHelper(c *gin.Context) *relaymodel.ErrorWithStatusCode {
+	ctx := c.Request.Context()
+	meta := meta.GetByContext(c)
+	// get & validate textRequest
+	textRequest, err := getAndValidateTextRequest(c, meta.Mode)
+	if err != nil {
+		logger.Errorf(ctx, "getAndValidateTextRequest failed: %s", err.Error())
+		return openai.ErrorWrapper(err, "invalid_text_request", http.StatusBadRequest)
+	}
+	meta.IsStream = textRequest.Stream
+
+	// map model name
+	meta.OriginModelName = textRequest.Model
+	textRequest.Model, _ = getMappedModelName(textRequest.Model, meta.ModelMapping)
+	meta.ActualModelName = textRequest.Model
+	// set system prompt if not empty
+	systemPromptReset := setSystemPrompt(ctx, textRequest, meta.SystemPrompt)
+	// get model ratio & group ratio
+	modelRatio := billingratio.GetModelRatio(textRequest.Model, meta.ChannelType)
+	// groupRatio := billingratio.GetGroupRatio(meta.Group)
+	groupRatio := c.GetFloat64(ctxkey.ChannelRatio)
+
+	ratio := modelRatio * groupRatio
+	// pre-consume quota
+	promptTokens := getPromptTokens(c.Request.Context(), textRequest, meta.Mode)
+	meta.PromptTokens = promptTokens
+	preConsumedQuota, bizErr := preConsumeQuota(c, textRequest, promptTokens, ratio, meta)
+	if bizErr != nil {
+		logger.Warnf(ctx, "preConsumeQuota failed: %+v", *bizErr)
+		return bizErr
+	}
+
+	adaptor := relay.GetAdaptor(meta.APIType)
+	if adaptor == nil {
+		return openai.ErrorWrapper(errors.Errorf("invalid api type: %d", meta.APIType), "invalid_api_type", http.StatusBadRequest)
+	}
+	adaptor.Init(meta)
+
+	// get request body
+	requestBody, err := getRequestBody(c, meta, textRequest, adaptor)
+	if err != nil {
+		return openai.ErrorWrapper(err, "convert_request_failed", http.StatusInternalServerError)
+	}
+
+	// for debug
+	requestBodyBytes, _ := io.ReadAll(requestBody)
+	requestBody = bytes.NewBuffer(requestBodyBytes)
+
+	// do request
+	resp, err := adaptor.DoRequest(c, meta, requestBody)
+	if err != nil {
+		logger.Errorf(ctx, "DoRequest failed: %s", err.Error())
+		return openai.ErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
+	}
+	if isErrorHappened(meta, resp) {
+		billing.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId)
+		return RelayErrorHandler(resp)
+	}
+
+	// do response
+	usage, respErr := adaptor.DoResponse(c, resp, meta)
+	if respErr != nil {
+		logger.Errorf(ctx, "respErr is not nil: %+v", respErr)
+		billing.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId)
+		return respErr
+	}
+
+	// post-consume quota
+	quotaId := c.GetInt(ctxkey.Id)
+	requestId := c.GetString(ctxkey.RequestId)
+	go func() {
+		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+		defer cancel()
+
+		quota := postConsumeQuota(ctx, usage, meta, textRequest, ratio, preConsumedQuota, modelRatio, groupRatio, systemPromptReset)
+
+		// also update user request cost
+		if quota != 0 {
+			docu := model.NewUserRequestCost(
+				quotaId,
+				requestId,
+				quota,
+			)
+			if err = docu.Insert(); err != nil {
+				logger.Errorf(ctx, "insert user request cost failed: %+v", err)
+			}
+		}
+	}()
+
+	return nil
+}
+
+func getRequestBody(c *gin.Context, meta *meta.Meta, textRequest *relaymodel.GeneralOpenAIRequest, adaptor adaptor.Adaptor) (io.Reader, error) {
+	if !config.EnforceIncludeUsage &&
+		meta.APIType == apitype.OpenAI &&
+		meta.OriginModelName == meta.ActualModelName &&
+		meta.ChannelType != channeltype.OpenAI && // openai also need to convert request
+		meta.ChannelType != channeltype.Baichuan {
+		return c.Request.Body, nil
+	}
+
+	// get request body
+	var requestBody io.Reader
+	convertedRequest, err := adaptor.ConvertRequest(c, meta.Mode, textRequest)
+	if err != nil {
+		logger.Debugf(c.Request.Context(), "converted request failed: %s\n", err.Error())
+		return nil, err
+	}
+	jsonData, err := json.Marshal(convertedRequest)
+	if err != nil {
+		logger.Debugf(c.Request.Context(), "converted request json_marshal_failed: %s\n", err.Error())
+		return nil, err
+	}
+	logger.Debugf(c.Request.Context(), "converted request: \n%s", string(jsonData))
+	requestBody = bytes.NewBuffer(jsonData)
+	return requestBody, nil
+}

relay/controller/validator/validation.go

@@ -0,0 +1,38 @@
+package validator
+
+import (
+	"math"
+
+	"github.com/pkg/errors"
+	"github.com/songquanpeng/one-api/relay/model"
+	"github.com/songquanpeng/one-api/relay/relaymode"
+)
+
+func ValidateTextRequest(textRequest *model.GeneralOpenAIRequest, relayMode int) error {
+	if textRequest.MaxTokens < 0 || textRequest.MaxTokens > math.MaxInt32/2 {
+		return errors.New("max_tokens is invalid")
+	}
+	if textRequest.Model == "" {
+		return errors.New("model is required")
+	}
+	switch relayMode {
+	case relaymode.Completions:
+		if textRequest.Prompt == "" {
+			return errors.New("field prompt is required")
+		}
+	case relaymode.ChatCompletions:
+		if textRequest.Messages == nil || len(textRequest.Messages) == 0 {
+			return errors.New("field messages is required")
+		}
+	case relaymode.Embeddings:
+	case relaymode.Moderations:
+		if textRequest.Input == "" {
+			return errors.New("field input is required")
+		}
+	case relaymode.Edits:
+		if textRequest.Instruction == "" {
+			return errors.New("field instruction is required")
+		}
+	}
+	return nil
+}