xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2026-01-01 12:35:58 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

♻️ refactor: 重构speech接口

Browse Source
This commit is contained in:
Martial BE
2023-11-29 18:11:15 +08:00
parent 1c7c2d40bb
commit 96dc7614e6
6 changed files with 125 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
controller/relay-text.go → controller/relay-helper.go
View File

@@ -13,7 +13,7 @@ import (
"github.com/gin-gonic/gin"
)
func relayTextHelper(c *gin.Context, relayMode int) *types.OpenAIErrorWithStatusCode {
func relayHelper(c *gin.Context, relayMode int) *types.OpenAIErrorWithStatusCode {
// 获取请求参数
channelType := c.GetInt("channel")
channelId := c.GetInt("channel_id")
@@ -36,10 +36,9 @@ func relayTextHelper(c *gin.Context, relayMode int) *types.OpenAIErrorWithStatus
return types.ErrorWrapper(err, "unmarshal_model_mapping_failed", http.StatusInternalServerError)
}
var promptTokens int
quotaInfo := &QuotaInfo{
modelName: "",
promptTokens: promptTokens,
promptTokens: 0,
userId: userId,
channelId: channelId,
tokenId: tokenId,
@@ -57,6 +56,8 @@ func relayTextHelper(c *gin.Context, relayMode int) *types.OpenAIErrorWithStatus
usage, openAIErrorWithStatusCode = handleEmbeddings(c, provider, modelMap, quotaInfo, group)
case common.RelayModeModerations:
usage, openAIErrorWithStatusCode = handleModerations(c, provider, modelMap, quotaInfo, group)
case common.RelayModeAudioSpeech:
usage, openAIErrorWithStatusCode = handleSpeech(c, provider, modelMap, quotaInfo, group)
default:
return types.ErrorWrapper(errors.New("invalid relay mode"), "invalid_relay_mode", http.StatusBadRequest)
}
@@ -112,6 +113,7 @@ func handleChatCompletions(c *gin.Context, provider providers_base.ProviderInter
promptTokens := common.CountTokenMessages(chatRequest.Messages, chatRequest.Model)
quotaInfo.modelName = chatRequest.Model
quotaInfo.promptTokens = promptTokens
quotaInfo.initQuotaInfo(group)
quota_err := quotaInfo.preQuotaConsumption()
if quota_err != nil {
@@ -144,6 +146,7 @@ func handleCompletions(c *gin.Context, provider providers_base.ProviderInterface
promptTokens := common.CountTokenInput(completionRequest.Prompt, completionRequest.Model)
quotaInfo.modelName = completionRequest.Model
quotaInfo.promptTokens = promptTokens
quotaInfo.initQuotaInfo(group)
quota_err := quotaInfo.preQuotaConsumption()
if quota_err != nil {
@@ -176,6 +179,7 @@ func handleEmbeddings(c *gin.Context, provider providers_base.ProviderInterface,
promptTokens := common.CountTokenInput(embeddingsRequest.Input, embeddingsRequest.Model)
quotaInfo.modelName = embeddingsRequest.Model
quotaInfo.promptTokens = promptTokens
quotaInfo.initQuotaInfo(group)
quota_err := quotaInfo.preQuotaConsumption()
if quota_err != nil {
@@ -212,6 +216,7 @@ func handleModerations(c *gin.Context, provider providers_base.ProviderInterface
promptTokens := common.CountTokenInput(moderationRequest.Input, moderationRequest.Model)
quotaInfo.modelName = moderationRequest.Model
quotaInfo.promptTokens = promptTokens
quotaInfo.initQuotaInfo(group)
quota_err := quotaInfo.preQuotaConsumption()
if quota_err != nil {
@@ -219,3 +224,36 @@ func handleModerations(c *gin.Context, provider providers_base.ProviderInterface
}
return moderationProvider.ModerationAction(&moderationRequest, isModelMapped, promptTokens)
}
func handleSpeech(c *gin.Context, provider providers_base.ProviderInterface, modelMap map[string]string, quotaInfo *QuotaInfo, group string) (*types.Usage, *types.OpenAIErrorWithStatusCode) {
var speechRequest types.SpeechAudioRequest
isModelMapped := false
speechProvider, ok := provider.(providers_base.SpeechInterface)
if !ok {
return nil, types.ErrorWrapper(errors.New("channel not implemented"), "channel_not_implemented", http.StatusNotImplemented)
}
err := common.UnmarshalBodyReusable(c, &speechRequest)
if err != nil {
return nil, types.ErrorWrapper(err, "bind_request_body_failed", http.StatusBadRequest)
}
if speechRequest.Input == "" {
return nil, types.ErrorWrapper(errors.New("field input is required"), "required_field_missing", http.StatusBadRequest)
}
if modelMap != nil && modelMap[speechRequest.Model] != "" {
speechRequest.Model = modelMap[speechRequest.Model]
isModelMapped = true
}
promptTokens := len(speechRequest.Input)
quotaInfo.modelName = speechRequest.Model
quotaInfo.promptTokens = promptTokens
quotaInfo.initQuotaInfo(group)
quota_err := quotaInfo.preQuotaConsumption()
if quota_err != nil {
return nil, quota_err
}
return speechProvider.SpeechAction(&speechRequest, isModelMapped, promptTokens)
}

6
controller/relay.go
View File

@@ -235,13 +235,13 @@ func Relay(c *gin.Context) {
relayMode = common.RelayModeEmbeddings
} else if strings.HasPrefix(c.Request.URL.Path, "/v1/moderations") {
relayMode = common.RelayModeModerations
} else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/speech") {
relayMode = common.RelayModeAudioSpeech
}
// } else if strings.HasPrefix(c.Request.URL.Path, "/v1/images/generations") {
// relayMode = RelayModeImagesGenerations
// } else if strings.HasPrefix(c.Request.URL.Path, "/v1/edits") {
// relayMode = RelayModeEdits
// } else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/speech") {
// relayMode = RelayModeAudioSpeech
// } else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcriptions") {
// relayMode = RelayModeAudioTranscription
// } else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/translations") {
@@ -257,7 +257,7 @@ func Relay(c *gin.Context) {
// case RelayModeAudioTranscription:
// err = relayAudioHelper(c, relayMode)
default:
err = relayTextHelper(c, relayMode)
err = relayHelper(c, relayMode)
}
if err != nil {
requestId := c.GetString(common.RequestIdKey)