feat: support tts model (#713)

* Added support for Text-to-Speech models and
endpoints

* chore: update impl

---------

Co-authored-by: JustSong <songquanpeng@foxmail.com>

common/model-ratio.go

@@ -59,7 +59,11 @@ var ModelRatio = map[string]float64{
 	"text-davinci-003":          10,
 	"text-davinci-edit-001":     10,
 	"code-davinci-edit-001":     10,
-	"whisper-1":                 15, // $0.006 / minute -> $0.006 / 150 words -> $0.006 / 200 tokens -> $0.03 / 1k tokens
+	"whisper-1":                 15,  // $0.006 / minute -> $0.006 / 150 words -> $0.006 / 200 tokens -> $0.03 / 1k tokens
+	"tts-1":                     7.5, // $0.015 / 1K characters
+	"tts-1-1106":                7.5,
+	"tts-1-hd":                  15, // $0.030 / 1K characters
+	"tts-1-hd-1106":             15,
 	"davinci":                   10,
 	"curie":                     10,
 	"babbage":                   10,

controller/model.go

@@ -81,6 +81,42 @@ func init() {
 			Root:       "whisper-1",
 			Parent:     nil,
 		},
+		{
+			Id:         "tts-1",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "tts-1",
+			Parent:     nil,
+		},
+		{
+			Id:         "tts-1-1106",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "tts-1-1106",
+			Parent:     nil,
+		},
+		{
+			Id:         "tts-1-hd",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "tts-1-hd",
+			Parent:     nil,
+		},
+		{
+			Id:         "tts-1-hd-1106",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "tts-1-hd-1106",
+			Parent:     nil,
+		},
 		{
 			Id:         "gpt-3.5-turbo",
 			Object:     "model",

controller/relay-audio.go

@@ -5,7 +5,6 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
-	"fmt"
 	"github.com/gin-gonic/gin"
 	"io"
 	"net/http"
@@ -21,6 +20,22 @@ func relayAudioHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 	channelId := c.GetInt("channel_id")
 	userId := c.GetInt("id")
 	group := c.GetString("group")
+	tokenName := c.GetString("token_name")
+
+	var ttsRequest TextToSpeechRequest
+	if relayMode == RelayModeAudioSpeech {
+		// Read JSON
+		err := common.UnmarshalBodyReusable(c, &ttsRequest)
+		// Check if JSON is valid
+		if err != nil {
+			return errorWrapper(err, "invalid_json", http.StatusBadRequest)
+		}
+		audioModel = ttsRequest.Model
+		// Check if text is too long 4096
+		if len(ttsRequest.Input) > 4096 {
+			return errorWrapper(errors.New("input is too long (over 4096 characters)"), "text_too_long", http.StatusBadRequest)
+		}
+	}
 
 	preConsumedTokens := common.PreConsumedQuota
 	modelRatio := common.GetModelRatio(audioModel)
@@ -31,22 +46,32 @@ func relayAudioHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 	if err != nil {
 		return errorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError)
 	}
-	if userQuota-preConsumedQuota < 0 {
+
-		return errorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden)
+	quota := 0
-	}
+	// Check if user quota is enough
-	err = model.CacheDecreaseUserQuota(userId, preConsumedQuota)
+	if relayMode == RelayModeAudioSpeech {
-	if err != nil {
+		quota = int(float64(len(ttsRequest.Input)) * modelRatio * groupRatio)
-		return errorWrapper(err, "decrease_user_quota_failed", http.StatusInternalServerError)
+		if quota > userQuota {
-	}
+			return errorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden)
-	if userQuota > 100*preConsumedQuota {
+		}
-		// in this case, we do not pre-consume quota
+	} else {
-		// because the user has enough quota
+		if userQuota-preConsumedQuota < 0 {
-		preConsumedQuota = 0
+			return errorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden)
-	}
+		}
-	if preConsumedQuota > 0 {
+		err = model.CacheDecreaseUserQuota(userId, preConsumedQuota)
-		err := model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
 		if err != nil {
-			return errorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden)
+			return 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 errorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden)
+			}
 		}
 	}
 
@@ -93,47 +118,32 @@ func relayAudioHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 	if err != nil {
 		return errorWrapper(err, "close_request_body_failed", http.StatusInternalServerError)
 	}
-	var audioResponse AudioResponse
 
-	defer func(ctx context.Context) {
+	if relayMode == RelayModeAudioSpeech {
-		go func() {
+		defer func(ctx context.Context) {
-			quota := countTokenText(audioResponse.Text, audioModel)
+			go postConsumeQuota(ctx, tokenId, quota, userId, channelId, modelRatio, groupRatio, audioModel, tokenName)
+		}(c.Request.Context())
+	} else {
+		responseBody, err := io.ReadAll(resp.Body)
+		if err != nil {
+			return errorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
+		}
+		err = resp.Body.Close()
+		if err != nil {
+			return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError)
+		}
+		var whisperResponse WhisperResponse
+		err = json.Unmarshal(responseBody, &whisperResponse)
+		if err != nil {
+			return errorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError)
+		}
+		defer func(ctx context.Context) {
+			quota := countTokenText(whisperResponse.Text, audioModel)
 			quotaDelta := quota - preConsumedQuota
-			err := model.PostConsumeTokenQuota(tokenId, quotaDelta)
+			go postConsumeQuota(ctx, tokenId, quotaDelta, userId, channelId, modelRatio, groupRatio, audioModel, tokenName)
-			if err != nil {
+		}(c.Request.Context())
-				common.SysError("error consuming token remain quota: " + err.Error())
+		resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
-			}
-			err = model.CacheUpdateUserQuota(userId)
-			if err != nil {
-				common.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, audioModel, 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 errorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
 	}
-	err = resp.Body.Close()
-	if err != nil {
-		return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError)
-	}
-	err = json.Unmarshal(responseBody, &audioResponse)
-	if err != nil {
-		return 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])
 	}

controller/relay-utils.go

@@ -1,6 +1,7 @@
 package controller
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"github.com/gin-gonic/gin"
@@ -8,6 +9,7 @@ import (
 	"io"
 	"net/http"
 	"one-api/common"
+	"one-api/model"
 	"strconv"
 	"strings"
 )
@@ -186,3 +188,20 @@ func getFullRequestURL(baseURL string, requestURL string, channelType int) strin
 	}
 	return fullRequestURL
 }
+
+func postConsumeQuota(ctx context.Context, tokenId int, quota int, userId int, channelId int, modelRatio float64, groupRatio float64, modelName string, tokenName string) {
+	err := model.PostConsumeTokenQuota(tokenId, quota)
+	if err != nil {
+		common.SysError("error consuming token remain quota: " + err.Error())
+	}
+	err = model.CacheUpdateUserQuota(userId)
+	if err != nil {
+		common.SysError("error update user quota cache: " + err.Error())
+	}
+	if quota != 0 {
+		logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio)
+		model.RecordConsumeLog(ctx, userId, channelId, 0, 0, modelName, tokenName, quota, logContent)
+		model.UpdateUserUsedQuotaAndRequestCount(userId, quota)
+		model.UpdateChannelUsedQuota(channelId, quota)
+	}
+}

controller/relay.go

@@ -24,7 +24,9 @@ const (
 	RelayModeModerations
 	RelayModeImagesGenerations
 	RelayModeEdits
-	RelayModeAudio
+	RelayModeAudioSpeech
+	RelayModeAudioTranscription
+	RelayModeAudioTranslation
 )
 
 // https://platform.openai.com/docs/api-reference/chat
@@ -89,10 +91,18 @@ type ImageRequest struct {
 	User           string `json:"user"`
 }
 
-type AudioResponse struct {
+type WhisperResponse struct {
 	Text string `json:"text,omitempty"`
 }
 
+type TextToSpeechRequest struct {
+	Model          string  `json:"model" binding:"required"`
+	Input          string  `json:"input" binding:"required"`
+	Voice          string  `json:"voice" binding:"required"`
+	Speed          float64 `json:"speed"`
+	ResponseFormat string  `json:"response_format"`
+}
+
 type Usage struct {
 	PromptTokens     int `json:"prompt_tokens"`
 	CompletionTokens int `json:"completion_tokens"`
@@ -189,14 +199,22 @@ func Relay(c *gin.Context) {
 		relayMode = RelayModeImagesGenerations
 	} else if strings.HasPrefix(c.Request.URL.Path, "/v1/edits") {
 		relayMode = RelayModeEdits
-	} else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio") {
+	} else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/speech") {
-		relayMode = RelayModeAudio
+		relayMode = RelayModeAudioSpeech
+	} else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcription") {
+		relayMode = RelayModeAudioTranscription
+	} else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/translation") {
+		relayMode = RelayModeAudioTranslation
 	}
 	var err *OpenAIErrorWithStatusCode
 	switch relayMode {
 	case RelayModeImagesGenerations:
 		err = relayImageHelper(c, relayMode)
-	case RelayModeAudio:
+	case RelayModeAudioSpeech:
+		fallthrough
+	case RelayModeAudioTranslation:
+		fallthrough
+	case RelayModeAudioTranscription:
 		err = relayAudioHelper(c, relayMode)
 	default:
 		err = relayTextHelper(c, relayMode)

middleware/distributor.go

@@ -40,10 +40,7 @@ func Distribute() func(c *gin.Context) {
 		} else {
 			// Select a channel for the user
 			var modelRequest ModelRequest
-			var err error
+			err := common.UnmarshalBodyReusable(c, &modelRequest)
-			if !strings.HasPrefix(c.Request.URL.Path, "/v1/audio") {
-				err = common.UnmarshalBodyReusable(c, &modelRequest)
-			}
 			if err != nil {
 				abortWithMessage(c, http.StatusBadRequest, "无效的请求")
 				return
@@ -60,10 +57,10 @@ func Distribute() func(c *gin.Context) {
 			}
 			if strings.HasPrefix(c.Request.URL.Path, "/v1/images/generations") {
 				if modelRequest.Model == "" {
-					modelRequest.Model = "dall-e"
+					modelRequest.Model = "dall-e-2"
 				}
 			}
-			if strings.HasPrefix(c.Request.URL.Path, "/v1/audio") {
+			if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcriptions") || strings.HasPrefix(c.Request.URL.Path, "/v1/audio/translations") {
 				if modelRequest.Model == "" {
 					modelRequest.Model = "whisper-1"
 				}

router/relay-router.go

@@ -29,6 +29,7 @@ func SetRelayRouter(router *gin.Engine) {
 		relayV1Router.POST("/engines/:model/embeddings", controller.Relay)
 		relayV1Router.POST("/audio/transcriptions", controller.Relay)
 		relayV1Router.POST("/audio/translations", controller.Relay)
+		relayV1Router.POST("/audio/speech", controller.Relay)
 		relayV1Router.GET("/files", controller.RelayNotImplemented)
 		relayV1Router.POST("/files", controller.RelayNotImplemented)
 		relayV1Router.DELETE("/files/:id", controller.RelayNotImplemented)