ai-gateway/controller/relay-translations.go
2023-12-26 16:40:50 +08:00

90 lines
2.2 KiB
Go

package controller
import (
"context"
"net/http"
"one-api/common"
"one-api/model"
providersBase "one-api/providers/base"
"one-api/types"
"github.com/gin-gonic/gin"
)
func RelayTranslations(c *gin.Context) {
var audioRequest types.AudioRequest
if err := common.UnmarshalBodyReusable(c, &audioRequest); err != nil {
common.AbortWithMessage(c, http.StatusBadRequest, err.Error())
return
}
channel, pass := fetchChannel(c, audioRequest.Model)
if pass {
return
}
// 解析模型映射
var isModelMapped bool
modelMap, err := parseModelMapping(channel.GetModelMapping())
if err != nil {
common.AbortWithMessage(c, http.StatusInternalServerError, err.Error())
return
}
if modelMap != nil && modelMap[audioRequest.Model] != "" {
audioRequest.Model = modelMap[audioRequest.Model]
isModelMapped = true
}
// 获取供应商
provider, pass := getProvider(c, channel, common.RelayModeAudioTranslation)
if pass {
return
}
translationProvider, ok := provider.(providersBase.TranslationInterface)
if !ok {
common.AbortWithMessage(c, http.StatusNotImplemented, "channel not implemented")
return
}
// 获取Input Tokens
promptTokens := 0
var quotaInfo *QuotaInfo
var errWithCode *types.OpenAIErrorWithStatusCode
var usage *types.Usage
quotaInfo, errWithCode = generateQuotaInfo(c, audioRequest.Model, promptTokens)
if errWithCode != nil {
errorHelper(c, errWithCode)
return
}
usage, errWithCode = translationProvider.TranslationAction(&audioRequest, isModelMapped, promptTokens)
// 如果报错,则退还配额
if errWithCode != nil {
tokenId := c.GetInt("token_id")
if quotaInfo.HandelStatus {
go func(ctx context.Context) {
// return pre-consumed quota
err := model.PostConsumeTokenQuota(tokenId, -quotaInfo.preConsumedQuota)
if err != nil {
common.LogError(ctx, "error return pre-consumed quota: "+err.Error())
}
}(c.Request.Context())
}
errorHelper(c, errWithCode)
return
} else {
tokenName := c.GetString("token_name")
// 如果没有报错,则消费配额
go func(ctx context.Context) {
err = quotaInfo.completedQuotaConsumption(usage, tokenName, ctx)
if err != nil {
common.LogError(ctx, err.Error())
}
}(c.Request.Context())
}
}