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

refactor: 重构计费代码

Browse Source
This commit is contained in:
CaIon 2024-04-23 23:51:27 +08:00
parent 2841669246
commit 783e8fd74a
1 changed files with 15 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
service/token_counter.go
View File

@ -128,7 +128,7 @@ func CountTokenChatRequest(request dto.GeneralOpenAIRequest, model string, check
var openaiTools []dto.OpenAITools var openaiTools []dto.OpenAITools
err := json.Unmarshal(toolsData, &openaiTools) err := json.Unmarshal(toolsData, &openaiTools)
if err != nil { if err != nil {
return 0, errors.New(fmt.Sprintf("count tools token fail: %s", err.Error())), false return 0, errors.New(fmt.Sprintf("count_tools_token_fail: %s", err.Error())), false
} }
countStr := "" countStr := ""
for _, tool := range openaiTools { for _, tool := range openaiTools {
@ -173,26 +173,23 @@ func CountTokenMessages(messages []dto.Message, model string, checkSensitive boo
tokenNum += tokensPerMessage tokenNum += tokensPerMessage
tokenNum += getTokenNum(tokenEncoder, message.Role) tokenNum += getTokenNum(tokenEncoder, message.Role)
if len(message.Content) > 0 { if len(message.Content) > 0 {
var arrayContent []dto.MediaMessage if message.IsStringContent() {
if err := json.Unmarshal(message.Content, &arrayContent); err != nil { stringContent := message.StringContent()
var stringContent string if checkSensitive {
if err := json.Unmarshal(message.Content, &stringContent); err != nil { contains, words := SensitiveWordContains(stringContent)
return 0, err, false if contains {
} else { err := fmt.Errorf("message contains sensitive words: [%s]", strings.Join(words, ", "))
if checkSensitive { return 0, err, true
contains, words := SensitiveWordContains(stringContent)
if contains {
err := fmt.Errorf("message contains sensitive words: [%s]", strings.Join(words, ", "))
return 0, err, true
}
}
tokenNum += getTokenNum(tokenEncoder, stringContent)
if message.Name != nil {
tokenNum += tokensPerName
tokenNum += getTokenNum(tokenEncoder, *message.Name)
} }
} }
tokenNum += getTokenNum(tokenEncoder, stringContent)
if message.Name != nil {
tokenNum += tokensPerName
tokenNum += getTokenNum(tokenEncoder, *message.Name)
}
} else { } else {
var err error
arrayContent := message.ParseContent()
for _, m := range arrayContent { for _, m := range arrayContent {
if m.Type == "image_url" { if m.Type == "image_url" {
var imageTokenNum int var imageTokenNum int