CaIon 1 год назад
Родитель
Сommit
783e8fd74a
1 измененных файлов с 15 добавлено и 18 удалено
  1. 15 18
      service/token_counter.go

+ 15 - 18
service/token_counter.go

@@ -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 err := json.Unmarshal(message.Content, &arrayContent); err != nil {
-				var stringContent string
-				if err := json.Unmarshal(message.Content, &stringContent); err != nil {
-					return 0, err, false
-				} else {
-					if checkSensitive {
-						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)
+			if message.IsStringContent() {
+				stringContent := message.StringContent()
+				if checkSensitive {
+					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)
+				}
 			} 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