Merge remote-tracking branch 'origin/upstream/main'

relay/controller/audio.go

@@ -0,0 +1,267 @@
+package controller
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/songquanpeng/one-api/common"
+	"github.com/songquanpeng/one-api/common/config"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/model"
+	"github.com/songquanpeng/one-api/relay/channel/openai"
+	"github.com/songquanpeng/one-api/relay/constant"
+	"github.com/songquanpeng/one-api/relay/util"
+	"io"
+	"net/http"
+	"strings"
+)
+
+func RelayAudioHelper(c *gin.Context, relayMode int) *openai.ErrorWithStatusCode {
+	audioModel := "whisper-1"
+
+	tokenId := c.GetInt("token_id")
+	channelType := c.GetInt("channel")
+	channelId := c.GetInt("channel_id")
+	userId := c.GetInt("id")
+	group := c.GetString("group")
+	tokenName := c.GetString("token_name")
+
+	var ttsRequest openai.TextToSpeechRequest
+	if relayMode == constant.RelayModeAudioSpeech {
+		// 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 := common.GetModelRatio(audioModel)
+	groupRatio := common.GetGroupRatio(group)
+	ratio := modelRatio * groupRatio
+	var quota int
+	var preConsumedQuota int
+	switch relayMode {
+	case constant.RelayModeAudioSpeech:
+		preConsumedQuota = int(float64(len(ttsRequest.Input)) * ratio)
+		quota = preConsumedQuota
+	default:
+		preConsumedQuota = int(float64(config.PreConsumedQuota) * ratio)
+	}
+	userQuota, err := model.CacheGetUserQuota(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 {
+		// 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)
+		}
+	}
+
+	// map model name
+	modelMapping := c.GetString("model_mapping")
+	if modelMapping != "" {
+		modelMap := make(map[string]string)
+		err := json.Unmarshal([]byte(modelMapping), &modelMap)
+		if err != nil {
+			return openai.ErrorWrapper(err, "unmarshal_model_mapping_failed", http.StatusInternalServerError)
+		}
+		if modelMap[audioModel] != "" {
+			audioModel = modelMap[audioModel]
+		}
+	}
+
+	baseURL := common.ChannelBaseURLs[channelType]
+	requestURL := c.Request.URL.String()
+	if c.GetString("base_url") != "" {
+		baseURL = c.GetString("base_url")
+	}
+
+	fullRequestURL := util.GetFullRequestURL(baseURL, requestURL, channelType)
+	if relayMode == constant.RelayModeAudioTranscription && channelType == common.ChannelTypeAzure {
+		// https://learn.microsoft.com/en-us/azure/ai-services/openai/whisper-quickstart?tabs=command-line#rest-api
+		apiVersion := util.GetAzureAPIVersion(c)
+		fullRequestURL = fmt.Sprintf("%s/openai/deployments/%s/audio/transcriptions?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 == constant.RelayModeAudioTranscription && channelType == common.ChannelTypeAzure {
+		// 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 := util.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)
+	}
+
+	if relayMode != constant.RelayModeAudioSpeech {
+		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(fmt.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 = openai.CountTokenText(text, audioModel)
+		resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
+	}
+	if resp.StatusCode != http.StatusOK {
+		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())
+		}
+		return util.RelayErrorHandler(resp)
+	}
+	quotaDelta := quota - preConsumedQuota
+	defer func(ctx context.Context) {
+		go util.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 "", fmt.Errorf("unmarshal_response_body_failed err :%w", err)
+	}
+	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 "", fmt.Errorf("unmarshal_response_body_failed err :%w", err)
+	}
+	return whisperResponse.Text, nil
+}

relay/controller/helper.go

@@ -0,0 +1,146 @@
+package controller
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/songquanpeng/one-api/common"
+	"github.com/songquanpeng/one-api/common/config"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/model"
+	"github.com/songquanpeng/one-api/relay/channel/openai"
+	"github.com/songquanpeng/one-api/relay/constant"
+	"github.com/songquanpeng/one-api/relay/util"
+	"io"
+	"math"
+	"net/http"
+)
+
+func getAndValidateTextRequest(c *gin.Context, relayMode int) (*openai.GeneralOpenAIRequest, error) {
+	textRequest := &openai.GeneralOpenAIRequest{}
+	err := common.UnmarshalBodyReusable(c, textRequest)
+	if err != nil {
+		return nil, err
+	}
+	if relayMode == constant.RelayModeModerations && textRequest.Model == "" {
+		textRequest.Model = "text-moderation-latest"
+	}
+	if relayMode == constant.RelayModeEmbeddings && textRequest.Model == "" {
+		textRequest.Model = c.Param("model")
+	}
+	err = util.ValidateTextRequest(textRequest, relayMode)
+	if err != nil {
+		return nil, err
+	}
+	return textRequest, nil
+}
+
+func getPromptTokens(textRequest *openai.GeneralOpenAIRequest, relayMode int) int {
+	switch relayMode {
+	case constant.RelayModeChatCompletions:
+		return openai.CountTokenMessages(textRequest.Messages, textRequest.Model)
+	case constant.RelayModeCompletions:
+		return openai.CountTokenInput(textRequest.Prompt, textRequest.Model)
+	case constant.RelayModeModerations:
+		return openai.CountTokenInput(textRequest.Input, textRequest.Model)
+	}
+	return 0
+}
+
+func getPreConsumedQuota(textRequest *openai.GeneralOpenAIRequest, promptTokens int, ratio float64) int {
+	preConsumedTokens := config.PreConsumedQuota
+	if textRequest.MaxTokens != 0 {
+		preConsumedTokens = promptTokens + textRequest.MaxTokens
+	}
+	return int(float64(preConsumedTokens) * ratio)
+}
+
+func preConsumeQuota(ctx context.Context, textRequest *openai.GeneralOpenAIRequest, promptTokens int, ratio float64, meta *util.RelayMeta) (int, *openai.ErrorWithStatusCode) {
+	preConsumedQuota := getPreConsumedQuota(textRequest, promptTokens, ratio)
+
+	userQuota, err := model.CacheGetUserQuota(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 {
+		// in this case, we do not pre-consume quota
+		// because the user has enough quota
+		preConsumedQuota = 0
+		logger.Info(ctx, 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 *openai.Usage, meta *util.RelayMeta, textRequest *openai.GeneralOpenAIRequest, ratio float64, preConsumedQuota int, modelRatio float64, groupRatio float64) {
+	if usage == nil {
+		logger.Error(ctx, "usage is nil, which is unexpected")
+		return
+	}
+	quota := 0
+	completionRatio := common.GetCompletionRatio(textRequest.Model)
+	promptTokens := usage.PromptTokens
+	completionTokens := usage.CompletionTokens
+	quota = int(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(meta.UserId)
+	if err != nil {
+		logger.Error(ctx, "error update user quota cache: "+err.Error())
+	}
+	if quota != 0 {
+		logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f，补全倍率 %.2f", modelRatio, groupRatio, completionRatio)
+		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)
+	}
+}
+
+func doRequest(ctx context.Context, c *gin.Context, meta *util.RelayMeta, isStream bool, fullRequestURL string, requestBody io.Reader) (*http.Response, error) {
+	req, err := http.NewRequest(c.Request.Method, fullRequestURL, requestBody)
+	if err != nil {
+		return nil, err
+	}
+	SetupRequestHeaders(c, req, meta, isStream)
+	resp, err := util.HTTPClient.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	if resp == nil {
+		return nil, errors.New("resp is nil")
+	}
+	err = req.Body.Close()
+	if err != nil {
+		logger.Warnf(ctx, "close req.Body failed: %+v", err)
+	}
+	err = c.Request.Body.Close()
+	if err != nil {
+		logger.Warnf(ctx, "close c.Request.Body failed: %+v", err)
+	}
+	return resp, nil
+}

relay/controller/image.go

@@ -0,0 +1,225 @@
+package controller
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"github.com/songquanpeng/one-api/common"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/model"
+	"github.com/songquanpeng/one-api/relay/channel/openai"
+	"github.com/songquanpeng/one-api/relay/util"
+	"io"
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+)
+
+func isWithinRange(element string, value int) bool {
+	if _, ok := common.DalleGenerationImageAmounts[element]; !ok {
+		return false
+	}
+	min := common.DalleGenerationImageAmounts[element][0]
+	max := common.DalleGenerationImageAmounts[element][1]
+
+	return value >= min && value <= max
+}
+
+func RelayImageHelper(c *gin.Context, relayMode int) *openai.ErrorWithStatusCode {
+	imageModel := "dall-e-2"
+	imageSize := "1024x1024"
+
+	tokenId := c.GetInt("token_id")
+	channelType := c.GetInt("channel")
+	channelId := c.GetInt("channel_id")
+	userId := c.GetInt("id")
+	group := c.GetString("group")
+
+	var imageRequest openai.ImageRequest
+	err := common.UnmarshalBodyReusable(c, &imageRequest)
+	if err != nil {
+		return openai.ErrorWrapper(err, "bind_request_body_failed", http.StatusBadRequest)
+	}
+
+	if imageRequest.N == 0 {
+		imageRequest.N = 1
+	}
+
+	// Size validation
+	if imageRequest.Size != "" {
+		imageSize = imageRequest.Size
+	}
+
+	// Model validation
+	if imageRequest.Model != "" {
+		imageModel = imageRequest.Model
+	}
+
+	imageCostRatio, hasValidSize := common.DalleSizeRatios[imageModel][imageSize]
+
+	// Check if model is supported
+	if hasValidSize {
+		if imageRequest.Quality == "hd" && imageModel == "dall-e-3" {
+			if imageSize == "1024x1024" {
+				imageCostRatio *= 2
+			} else {
+				imageCostRatio *= 1.5
+			}
+		}
+	} else {
+		return openai.ErrorWrapper(errors.New("size not supported for this image model"), "size_not_supported", http.StatusBadRequest)
+	}
+
+	// Prompt validation
+	if imageRequest.Prompt == "" {
+		return openai.ErrorWrapper(errors.New("prompt is required"), "prompt_missing", http.StatusBadRequest)
+	}
+
+	// Check prompt length
+	if len(imageRequest.Prompt) > common.DalleImagePromptLengthLimitations[imageModel] {
+		return openai.ErrorWrapper(errors.New("prompt is too long"), "prompt_too_long", http.StatusBadRequest)
+	}
+
+	// Number of generated images validation
+	if isWithinRange(imageModel, imageRequest.N) == false {
+		// channel not azure
+		if channelType != common.ChannelTypeAzure {
+			return openai.ErrorWrapper(errors.New("invalid value of n"), "n_not_within_range", http.StatusBadRequest)
+		}
+	}
+
+	// map model name
+	modelMapping := c.GetString("model_mapping")
+	isModelMapped := false
+	if modelMapping != "" {
+		modelMap := make(map[string]string)
+		err := json.Unmarshal([]byte(modelMapping), &modelMap)
+		if err != nil {
+			return openai.ErrorWrapper(err, "unmarshal_model_mapping_failed", http.StatusInternalServerError)
+		}
+		if modelMap[imageModel] != "" {
+			imageModel = modelMap[imageModel]
+			isModelMapped = true
+		}
+	}
+	baseURL := common.ChannelBaseURLs[channelType]
+	requestURL := c.Request.URL.String()
+	if c.GetString("base_url") != "" {
+		baseURL = c.GetString("base_url")
+	}
+	fullRequestURL := util.GetFullRequestURL(baseURL, requestURL, channelType)
+	if channelType == common.ChannelTypeAzure {
+		// https://learn.microsoft.com/en-us/azure/ai-services/openai/dall-e-quickstart?tabs=dalle3%2Ccommand-line&pivots=rest-api
+		apiVersion := util.GetAzureAPIVersion(c)
+		// https://{resource_name}.openai.azure.com/openai/deployments/dall-e-3/images/generations?api-version=2023-06-01-preview
+		fullRequestURL = fmt.Sprintf("%s/openai/deployments/%s/images/generations?api-version=%s", baseURL, imageModel, apiVersion)
+	}
+
+	var requestBody io.Reader
+	if isModelMapped || channelType == common.ChannelTypeAzure { // make Azure channel request body
+		jsonStr, err := json.Marshal(imageRequest)
+		if err != nil {
+			return openai.ErrorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError)
+		}
+		requestBody = bytes.NewBuffer(jsonStr)
+	} else {
+		requestBody = c.Request.Body
+	}
+
+	modelRatio := common.GetModelRatio(imageModel)
+	groupRatio := common.GetGroupRatio(group)
+	ratio := modelRatio * groupRatio
+	userQuota, err := model.CacheGetUserQuota(userId)
+
+	quota := int(ratio*imageCostRatio*1000) * imageRequest.N
+
+	if userQuota-quota < 0 {
+		return openai.ErrorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden)
+	}
+
+	req, err := http.NewRequest(c.Request.Method, fullRequestURL, requestBody)
+	if err != nil {
+		return openai.ErrorWrapper(err, "new_request_failed", http.StatusInternalServerError)
+	}
+	token := c.Request.Header.Get("Authorization")
+	if channelType == common.ChannelTypeAzure { // Azure authentication
+		token = strings.TrimPrefix(token, "Bearer ")
+		req.Header.Set("api-key", token)
+	} else {
+		req.Header.Set("Authorization", token)
+	}
+
+	req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
+	req.Header.Set("Accept", c.Request.Header.Get("Accept"))
+
+	resp, err := util.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)
+	}
+	var textResponse openai.ImageResponse
+
+	defer func(ctx context.Context) {
+		if resp.StatusCode != http.StatusOK {
+			return
+		}
+		err := model.PostConsumeTokenQuota(tokenId, quota)
+		if err != nil {
+			logger.SysError("error consuming token remain quota: " + err.Error())
+		}
+		err = model.CacheUpdateUserQuota(userId)
+		if err != nil {
+			logger.SysError("error update user quota cache: " + err.Error())
+		}
+		if quota != 0 {
+			tokenName := c.GetString("token_name")
+			logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio)
+			model.RecordConsumeLog(ctx, userId, channelId, 0, 0, imageModel, tokenName, quota, logContent)
+			model.UpdateUserUsedQuotaAndRequestCount(userId, quota)
+			channelId := c.GetInt("channel_id")
+			model.UpdateChannelUsedQuota(channelId, quota)
+		}
+	}(c.Request.Context())
+
+	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)
+	}
+	err = json.Unmarshal(responseBody, &textResponse)
+	if err != nil {
+		return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError)
+	}
+
+	resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
+
+	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
+}

relay/controller/temp.go

@@ -0,0 +1,333 @@
+package controller
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+	"github.com/songquanpeng/one-api/common"
+	"github.com/songquanpeng/one-api/common/helper"
+	"github.com/songquanpeng/one-api/relay/channel/aiproxy"
+	"github.com/songquanpeng/one-api/relay/channel/anthropic"
+	"github.com/songquanpeng/one-api/relay/channel/google"
+	"github.com/songquanpeng/one-api/relay/channel/openai"
+	"github.com/songquanpeng/one-api/relay/constant"
+	"github.com/songquanpeng/one-api/relay/util"
+)
+
+func GetRequestURL(requestURL string, meta *util.RelayMeta, textRequest *openai.GeneralOpenAIRequest) (string, error) {
+	fullRequestURL := util.GetFullRequestURL(meta.BaseURL, requestURL, meta.ChannelType)
+	switch meta.APIType {
+	case constant.APITypeOpenAI:
+		if meta.ChannelType == common.ChannelTypeAzure {
+			// https://learn.microsoft.com/en-us/azure/cognitive-services/openai/chatgpt-quickstart?pivots=rest-api&tabs=command-line#rest-api
+			requestURL := strings.Split(requestURL, "?")[0]
+			requestURL = fmt.Sprintf("%s?api-version=%s", requestURL, meta.APIVersion)
+			task := strings.TrimPrefix(requestURL, "/v1/")
+			model_ := textRequest.Model
+			model_ = strings.Replace(model_, ".", "", -1)
+			// https://github.com/songquanpeng/one-api/issues/67
+			model_ = strings.TrimSuffix(model_, "-0301")
+			model_ = strings.TrimSuffix(model_, "-0314")
+			model_ = strings.TrimSuffix(model_, "-0613")
+
+			requestURL = fmt.Sprintf("/openai/deployments/%s/%s", model_, task)
+			fullRequestURL = util.GetFullRequestURL(meta.BaseURL, requestURL, meta.ChannelType)
+		}
+	case constant.APITypeClaude:
+		fullRequestURL = fmt.Sprintf("%s/v1/complete", meta.BaseURL)
+	// case constant.APITypeBaidu:
+	// 	switch textRequest.Model {
+	// 	case "ERNIE-Bot":
+	// 		fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions"
+	// 	case "ERNIE-Bot-turbo":
+	// 		fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant"
+	// 	case "ERNIE-Bot-4":
+	// 		fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro"
+	// 	case "BLOOMZ-7B":
+	// 		fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/bloomz_7b1"
+	// 	case "Embedding-V1":
+	// 		fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/embedding-v1"
+	// 	}
+	// 	var accessToken string
+	// 	var err error
+	// 	if accessToken, err = baidu.GetAccessToken(meta.APIKey); err != nil {
+	// 		return "", fmt.Errorf("failed to get baidu access token: %w", err)
+	// 	}
+	// 	fullRequestURL += "?access_token=" + accessToken
+	case constant.APITypePaLM:
+		fullRequestURL = fmt.Sprintf("%s/v1beta2/models/chat-bison-001:generateMessage", meta.BaseURL)
+	case constant.APITypeGemini:
+		version := helper.AssignOrDefault(meta.APIVersion, "v1")
+		action := "generateContent"
+		if textRequest.Stream {
+			action = "streamGenerateContent"
+		}
+		fullRequestURL = fmt.Sprintf("%s/%s/models/%s:%s", meta.BaseURL, version, textRequest.Model, action)
+	case constant.APITypeZhipu:
+		method := "invoke"
+		if textRequest.Stream {
+			method = "sse-invoke"
+		}
+		fullRequestURL = fmt.Sprintf("https://open.bigmodel.cn/api/paas/v3/model-api/%s/%s", textRequest.Model, method)
+	case constant.APITypeAli:
+		fullRequestURL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
+		if meta.Mode == constant.RelayModeEmbeddings {
+			fullRequestURL = "https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding"
+		}
+	// case constant.APITypeTencent:
+	// 	fullRequestURL = "https://hunyuan.cloud.tencent.com/hyllm/v1/chat/completions"
+	case constant.APITypeAIProxyLibrary:
+		fullRequestURL = fmt.Sprintf("%s/api/library/ask", meta.BaseURL)
+	}
+	return fullRequestURL, nil
+}
+
+func GetRequestBody(c *gin.Context, textRequest openai.GeneralOpenAIRequest, isModelMapped bool, apiType int, relayMode int) (io.Reader, error) {
+	var requestBody io.Reader
+	if isModelMapped {
+		jsonStr, err := json.Marshal(textRequest)
+		if err != nil {
+			return nil, err
+		}
+		requestBody = bytes.NewBuffer(jsonStr)
+	} else {
+		requestBody = c.Request.Body
+	}
+	switch apiType {
+	case constant.APITypeClaude:
+		claudeRequest := anthropic.ConvertRequest(textRequest)
+		jsonStr, err := json.Marshal(claudeRequest)
+		if err != nil {
+			return nil, err
+		}
+		requestBody = bytes.NewBuffer(jsonStr)
+	// case constant.APITypeBaidu:
+	// 	var jsonData []byte
+	// 	var err error
+	// 	switch relayMode {
+	// 	case constant.RelayModeEmbeddings:
+	// 		baiduEmbeddingRequest := baidu.ConvertEmbeddingRequest(textRequest)
+	// 		jsonData, err = json.Marshal(baiduEmbeddingRequest)
+	// 	default:
+	// 		baiduRequest := baidu.ConvertRequest(textRequest)
+	// 		jsonData, err = json.Marshal(baiduRequest)
+	// 	}
+	// 	if err != nil {
+	// 		return nil, err
+	// 	}
+	// 	requestBody = bytes.NewBuffer(jsonData)
+	case constant.APITypePaLM:
+		palmRequest := google.ConvertPaLMRequest(textRequest)
+		jsonStr, err := json.Marshal(palmRequest)
+		if err != nil {
+			return nil, err
+		}
+		requestBody = bytes.NewBuffer(jsonStr)
+	case constant.APITypeGemini:
+		geminiChatRequest := google.ConvertGeminiRequest(textRequest)
+		jsonStr, err := json.Marshal(geminiChatRequest)
+		if err != nil {
+			return nil, err
+		}
+		requestBody = bytes.NewBuffer(jsonStr)
+	// case constant.APITypeZhipu:
+	// 	zhipuRequest := zhipu.ConvertRequest(textRequest)
+	// 	jsonStr, err := json.Marshal(zhipuRequest)
+	// 	if err != nil {
+	// 		return nil, err
+	// 	}
+	// 	requestBody = bytes.NewBuffer(jsonStr)
+	// case constant.APITypeAli:
+	// 	var jsonStr []byte
+	// 	var err error
+	// 	switch relayMode {
+	// 	case constant.RelayModeEmbeddings:
+	// 		aliEmbeddingRequest := ali.ConvertEmbeddingRequest(textRequest)
+	// 		jsonStr, err = json.Marshal(aliEmbeddingRequest)
+	// 	default:
+	// 		aliRequest := ali.ConvertRequest(textRequest)
+	// 		jsonStr, err = json.Marshal(aliRequest)
+	// 	}
+	// 	if err != nil {
+	// 		return nil, err
+	// 	}
+	// 	requestBody = bytes.NewBuffer(jsonStr)
+	// case constant.APITypeTencent:
+	// 	apiKey := c.Request.Header.Get("Authorization")
+	// 	apiKey = strings.TrimPrefix(apiKey, "Bearer ")
+	// 	appId, secretId, secretKey, err := tencent.ParseConfig(apiKey)
+	// 	if err != nil {
+	// 		return nil, err
+	// 	}
+	// 	tencentRequest := tencent.ConvertRequest(textRequest)
+	// 	tencentRequest.AppId = appId
+	// 	tencentRequest.SecretId = secretId
+	// 	jsonStr, err := json.Marshal(tencentRequest)
+	// 	if err != nil {
+	// 		return nil, err
+	// 	}
+	// 	sign := tencent.GetSign(*tencentRequest, secretKey)
+	// 	c.Request.Header.Set("Authorization", sign)
+	// 	requestBody = bytes.NewBuffer(jsonStr)
+	case constant.APITypeAIProxyLibrary:
+		aiProxyLibraryRequest := aiproxy.ConvertRequest(textRequest)
+		aiProxyLibraryRequest.LibraryId = c.GetString("library_id")
+		jsonStr, err := json.Marshal(aiProxyLibraryRequest)
+		if err != nil {
+			return nil, err
+		}
+		requestBody = bytes.NewBuffer(jsonStr)
+	}
+	return requestBody, nil
+}
+
+func SetupRequestHeaders(c *gin.Context, req *http.Request, meta *util.RelayMeta, isStream bool) {
+	SetupAuthHeaders(c, req, meta, isStream)
+	req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
+	req.Header.Set("Accept", c.Request.Header.Get("Accept"))
+	if isStream && c.Request.Header.Get("Accept") == "" {
+		req.Header.Set("Accept", "text/event-stream")
+	}
+}
+
+func SetupAuthHeaders(c *gin.Context, req *http.Request, meta *util.RelayMeta, isStream bool) {
+	apiKey := meta.APIKey
+	switch meta.APIType {
+	case constant.APITypeOpenAI:
+		if meta.ChannelType == common.ChannelTypeAzure {
+			req.Header.Set("api-key", apiKey)
+		} else {
+			req.Header.Set("Authorization", c.Request.Header.Get("Authorization"))
+			if meta.ChannelType == common.ChannelTypeOpenRouter {
+				req.Header.Set("HTTP-Referer", "https://github.com/songquanpeng/one-api")
+				req.Header.Set("X-Title", "One API")
+			}
+		}
+	case constant.APITypeClaude:
+		req.Header.Set("x-api-key", apiKey)
+		anthropicVersion := c.Request.Header.Get("anthropic-version")
+		if anthropicVersion == "" {
+			anthropicVersion = "2023-06-01"
+		}
+		req.Header.Set("anthropic-version", anthropicVersion)
+	// case constant.APITypeZhipu:
+	// 	token := zhipu.GetToken(apiKey)
+	// 	req.Header.Set("Authorization", token)
+	// case constant.APITypeAli:
+	// 	req.Header.Set("Authorization", "Bearer "+apiKey)
+	// 	if isStream {
+	// 		req.Header.Set("X-DashScope-SSE", "enable")
+	// 	}
+	// 	if c.GetString("plugin") != "" {
+	// 		req.Header.Set("X-DashScope-Plugin", c.GetString("plugin"))
+	// 	}
+	// case constant.APITypeTencent:
+	// 	req.Header.Set("Authorization", apiKey)
+	case constant.APITypePaLM:
+		req.Header.Set("x-goog-api-key", apiKey)
+	case constant.APITypeGemini:
+		req.Header.Set("x-goog-api-key", apiKey)
+	default:
+		req.Header.Set("Authorization", "Bearer "+apiKey)
+	}
+}
+
+func DoResponse(c *gin.Context, textRequest *openai.GeneralOpenAIRequest, resp *http.Response, relayMode int, apiType int, isStream bool, promptTokens int) (usage *openai.Usage, err *openai.ErrorWithStatusCode) {
+	var responseText string
+	switch apiType {
+	case constant.APITypeOpenAI:
+		if isStream {
+			err, responseText = openai.StreamHandler(c, resp, relayMode)
+		} else {
+			err, usage = openai.Handler(c, resp, promptTokens, textRequest.Model)
+		}
+	case constant.APITypeClaude:
+		if isStream {
+			err, responseText = anthropic.StreamHandler(c, resp)
+		} else {
+			err, usage = anthropic.Handler(c, resp, promptTokens, textRequest.Model)
+		}
+	// case constant.APITypeBaidu:
+	// 	if isStream {
+	// 		err, usage = baidu.StreamHandler(c, resp)
+	// 	} else {
+	// 		switch relayMode {
+	// 		case constant.RelayModeEmbeddings:
+	// 			err, usage = baidu.EmbeddingHandler(c, resp)
+	// 		default:
+	// 			err, usage = baidu.Handler(c, resp)
+	// 		}
+	// 	}
+	case constant.APITypePaLM:
+		if isStream { // PaLM2 API does not support stream
+			err, responseText = google.PaLMStreamHandler(c, resp)
+		} else {
+			err, usage = google.PaLMHandler(c, resp, promptTokens, textRequest.Model)
+		}
+	case constant.APITypeGemini:
+		if isStream {
+			err, responseText = google.StreamHandler(c, resp)
+		} else {
+			err, usage = google.GeminiHandler(c, resp, promptTokens, textRequest.Model)
+		}
+	// case constant.APITypeZhipu:
+	// 	if isStream {
+	// 		err, usage = zhipu.StreamHandler(c, resp)
+	// 	} else {
+	// 		err, usage = zhipu.Handler(c, resp)
+	// 	}
+	// case constant.APITypeAli:
+	// 	if isStream {
+	// 		err, usage = ali.StreamHandler(c, resp)
+	// 	} else {
+	// 		switch relayMode {
+	// 		case constant.RelayModeEmbeddings:
+	// 			err, usage = ali.EmbeddingHandler(c, resp)
+	// 		default:
+	// 			err, usage = ali.Handler(c, resp)
+	// 		}
+	// 	}
+	// case constant.APITypeXunfei:
+	// 	auth := c.Request.Header.Get("Authorization")
+	// 	auth = strings.TrimPrefix(auth, "Bearer ")
+	// 	splits := strings.Split(auth, "|")
+	// 	if len(splits) != 3 {
+	// 		return nil, openai.ErrorWrapper(errors.New("invalid auth"), "invalid_auth", http.StatusBadRequest)
+	// 	}
+	// 	if isStream {
+	// 		err, usage = xunfei.StreamHandler(c, *textRequest, splits[0], splits[1], splits[2])
+	// 	} else {
+	// 		err, usage = xunfei.Handler(c, *textRequest, splits[0], splits[1], splits[2])
+	// 	}
+	case constant.APITypeAIProxyLibrary:
+		if isStream {
+			err, usage = aiproxy.StreamHandler(c, resp)
+		} else {
+			err, usage = aiproxy.Handler(c, resp)
+		}
+	// case constant.APITypeTencent:
+	// 	if isStream {
+	// 		err, responseText = tencent.StreamHandler(c, resp)
+	// 	} else {
+	// 		err, usage = tencent.Handler(c, resp)
+	// 	}
+	default:
+		return nil, openai.ErrorWrapper(errors.New("unknown api type"), "unknown_api_type", http.StatusInternalServerError)
+	}
+	if err != nil {
+		return nil, err
+	}
+	if usage == nil && responseText != "" {
+		usage = &openai.Usage{}
+		usage.PromptTokens = promptTokens
+		usage.CompletionTokens = openai.CountTokenText(responseText, textRequest.Model)
+		usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
+	}
+	return usage, nil
+}

relay/controller/text.go

@@ -0,0 +1,76 @@
+package controller
+
+import (
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/songquanpeng/one-api/common"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/relay/channel/openai"
+	"github.com/songquanpeng/one-api/relay/constant"
+	"github.com/songquanpeng/one-api/relay/util"
+	"net/http"
+	"strings"
+)
+
+func RelayTextHelper(c *gin.Context) *openai.ErrorWithStatusCode {
+	ctx := c.Request.Context()
+	meta := util.GetRelayMeta(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)
+	}
+	// map model name
+	var isModelMapped bool
+	textRequest.Model, isModelMapped = util.GetMappedModelName(textRequest.Model, meta.ModelMapping)
+	// get model ratio & group ratio
+	modelRatio := common.GetModelRatio(textRequest.Model)
+	groupRatio := common.GetGroupRatio(meta.Group)
+	ratio := modelRatio * groupRatio
+	// pre-consume quota
+	promptTokens := getPromptTokens(textRequest, meta.Mode)
+	preConsumedQuota, bizErr := preConsumeQuota(ctx, textRequest, promptTokens, ratio, meta)
+	if bizErr != nil {
+		logger.Warnf(ctx, "preConsumeQuota failed: %+v", *bizErr)
+		return bizErr
+	}
+
+	// get request body
+	requestBody, err := GetRequestBody(c, *textRequest, isModelMapped, meta.APIType, meta.Mode)
+	if err != nil {
+		return openai.ErrorWrapper(err, "get_request_body_failed", http.StatusInternalServerError)
+	}
+	// do request
+	var resp *http.Response
+	isStream := textRequest.Stream
+	if meta.APIType != constant.APITypeXunfei { // cause xunfei use websocket
+		fullRequestURL, err := GetRequestURL(c.Request.URL.String(), meta, textRequest)
+		if err != nil {
+			logger.Error(ctx, fmt.Sprintf("util.GetRequestURL failed: %s", err.Error()))
+			return openai.ErrorWrapper(fmt.Errorf("util.GetRequestURL failed"), "get_request_url_failed", http.StatusInternalServerError)
+		}
+
+		resp, err = doRequest(ctx, c, meta, isStream, fullRequestURL, requestBody)
+		if err != nil {
+			logger.Errorf(ctx, "doRequest failed: %s", err.Error())
+			return openai.ErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
+		}
+		isStream = isStream || strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream")
+
+		if resp.StatusCode != http.StatusOK {
+			util.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId)
+			return util.RelayErrorHandler(resp)
+		}
+	}
+	// do response
+	usage, respErr := DoResponse(c, textRequest, resp, meta.Mode, meta.APIType, isStream, promptTokens)
+	if respErr != nil {
+		logger.Errorf(ctx, "respErr is not nil: %+v", respErr)
+		util.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId)
+		return respErr
+	}
+	// post-consume quota
+	go postConsumeQuota(ctx, usage, meta, textRequest, ratio, preConsumedQuota, modelRatio, groupRatio)
+	return nil
+}