xinyin025/new-api
1
0
Fork 0
mirror of https://github.com/linux-do/new-api.git synced 2025-11-09 23:53:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: 初步兼容敏感词过滤

Browse Source
This commit is contained in:
CaIon
2024-03-20 17:07:42 +08:00
parent bec21ade9d
commit 7a663d26ec
22 changed files with 293 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
service/usage_helpr.go
View File

@@ -1,27 +1,26 @@
package service
import (
"errors"
"one-api/dto"
"one-api/relay/constant"
)
func GetPromptTokens(textRequest dto.GeneralOpenAIRequest, relayMode int) (int, error) {
switch relayMode {
case constant.RelayModeChatCompletions:
return CountTokenMessages(textRequest.Messages, textRequest.Model)
case constant.RelayModeCompletions:
return CountTokenInput(textRequest.Prompt, textRequest.Model), nil
case constant.RelayModeModerations:
return CountTokenInput(textRequest.Input, textRequest.Model), nil
}
return 0, errors.New("unknown relay mode")
}
//func GetPromptTokens(textRequest dto.GeneralOpenAIRequest, relayMode int) (int, error) {
// switch relayMode {
// case constant.RelayModeChatCompletions:
// return CountTokenMessages(textRequest.Messages, textRequest.Model)
// case constant.RelayModeCompletions:
// return CountTokenInput(textRequest.Prompt, textRequest.Model), nil
// case constant.RelayModeModerations:
// return CountTokenInput(textRequest.Input, textRequest.Model), nil
// }
// return 0, errors.New("unknown relay mode")
//}
func ResponseText2Usage(responseText string, modeName string, promptTokens int) *dto.Usage {
func ResponseText2Usage(responseText string, modeName string, promptTokens int) (*dto.Usage, error) {
usage := &dto.Usage{}
usage.PromptTokens = promptTokens
usage.CompletionTokens = CountTokenText(responseText, modeName)
ctkm, err := CountTokenText(responseText, modeName, false)
usage.CompletionTokens = ctkm
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
return usage
return usage, err
}