token count update

Author: PlexPt

pom.xml

@@ -132,7 +132,7 @@
         <dependency>
             <groupId>com.knuddels</groupId>
             <artifactId>jtokkit</artifactId>
-            <version>0.4.0</version>
+            <version>1.0.0</version>
         </dependency>
     </dependencies>
 

src/main/java/com/plexpt/chatgpt/ChatGPTStream.java

@@ -1,22 +1,15 @@
 package com.plexpt.chatgpt;
 
-import com.fasterxml.jackson.databind.ObjectMapper;
+import cn.hutool.core.util.RandomUtil;
+import cn.hutool.http.ContentType;
 import com.plexpt.chatgpt.api.Api;
 import com.plexpt.chatgpt.entity.chat.ChatCompletion;
 import com.plexpt.chatgpt.entity.chat.Message;
-
-import java.net.Proxy;
-import java.util.List;
-import java.util.Objects;
-import java.util.concurrent.TimeUnit;
-
-import cn.hutool.core.util.RandomUtil;
-import cn.hutool.http.ContentType;
+import com.plexpt.chatgpt.util.fastjson.JSON;
 import lombok.AllArgsConstructor;
 import lombok.Builder;
 import lombok.Data;
 import lombok.NoArgsConstructor;
-import lombok.NonNull;
 import lombok.extern.slf4j.Slf4j;
 import okhttp3.MediaType;
 import okhttp3.OkHttpClient;
@@ -26,6 +19,11 @@
 import okhttp3.sse.EventSourceListener;
 import okhttp3.sse.EventSources;
 
+import java.net.Proxy;
+import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.TimeUnit;
+
 
 /**
  * open ai 客户端
@@ -89,8 +87,8 @@ public void streamChatCompletion(ChatCompletion chatCompletion,
 
         try {
             EventSource.Factory factory = EventSources.createFactory(okHttpClient);
-            ObjectMapper mapper = new ObjectMapper();
-            String requestBody = mapper.writeValueAsString(chatCompletion);
+
+            String requestBody = JSON.toJSONString(chatCompletion);
             String key = apiKey;
             if (apiKeyList != null && !apiKeyList.isEmpty()) {
                 key = RandomUtil.randomEle(apiKeyList);
@@ -99,8 +97,7 @@ public void streamChatCompletion(ChatCompletion chatCompletion,
 
             Request request = new Request.Builder()
                     .url(apiHost + "v1/chat/completions")
-                    .post(RequestBody.create(MediaType.parse(ContentType.JSON.getValue()),
-                            requestBody))
+                    .post(RequestBody.create(MediaType.parse(ContentType.JSON.getValue()), requestBody))
                     .header("Authorization", "Bearer " + key)
                     .build();
             factory.newEventSource(request, eventSourceListener);

src/main/java/com/plexpt/chatgpt/util/TokensUtil.java

@@ -1,64 +1,98 @@
 package com.plexpt.chatgpt.util;
 
-import cn.hutool.core.util.StrUtil;
 import com.knuddels.jtokkit.Encodings;
 import com.knuddels.jtokkit.api.Encoding;
 import com.knuddels.jtokkit.api.EncodingRegistry;
-import com.plexpt.chatgpt.entity.chat.ChatCompletion;
+import com.knuddels.jtokkit.api.EncodingType;
+import com.knuddels.jtokkit.api.ModelType;
 import com.plexpt.chatgpt.entity.chat.Message;
 import lombok.experimental.UtilityClass;
+import org.springframework.util.CollectionUtils;
+import org.springframework.util.StringUtils;
 
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 
 @UtilityClass
 public class TokensUtil {
 
-    private static final Map<String, Encoding> modelEncodingMap = new HashMap<>();
-    private static final EncodingRegistry encodingRegistry = Encodings.newDefaultEncodingRegistry();
+    public static EncodingRegistry registry = Encodings.newDefaultEncodingRegistry();
+    public static Encoding encoding = registry.getEncoding(EncodingType.CL100K_BASE);
 
-    static {
-        for (ChatCompletion.Model model : ChatCompletion.Model.values()) {
-            Optional<Encoding> encodingForModel = encodingRegistry.getEncodingForModel(model.getName());
-            encodingForModel.ifPresent(encoding -> modelEncodingMap.put(model.getName(), encoding));
-        }
+
+    /**
+     * 计算text信息的tokens
+     *
+     * @param text
+     * @return
+     */
+    public static int countTextTokens(String text) {
+        return encoding.countTokens(text);
     }
 
+
     /**
-     * 计算tokens
-     * @param modelName 模型名称
-     * @param messages 消息列表
-     * @return 计算出的tokens数量
+     * 获取modelType
+     *
+     * @param name
+     * @return
      */
+    private static ModelType getModelTypeByName(String name) {
+        Optional<ModelType> optional = ModelType.fromName(name);
 
-    public static int tokens(String modelName, List<Message> messages) {
-        Encoding encoding = modelEncodingMap.get(modelName);
-        if (encoding == null) {
-            throw new IllegalArgumentException("Unsupported model: " + modelName);
+        return optional.orElse(ModelType.GPT_3_5_TURBO);
+    }
+
+    /**
+     * 通过模型名称计算messages获取编码数组
+     * 参考官方的处理逻辑：
+     * <a href=https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb>https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb</a>
+     *
+     * @param messages 消息体
+     * @return
+     */
+    public static int tokens(List<Message> messages, String model) {
+        if (CollectionUtils.isEmpty(messages)) {
+            return 0;
         }
 
-        int tokensPerMessage = 0;
-        int tokensPerName = 0;
-        if (modelName.startsWith("gpt-4")) {
+        //"gpt-3.5-turbo"
+        // every message follows <|start|>{role/name}\n{content}<|end|>\n
+        int tokensPerMessage = 4;
+        // if there's a name, the role is omitted
+        int tokensPerName = -1;
+
+        if (StringUtils.startsWithIgnoreCase(model, ModelType.GPT_4.getName())) {
             tokensPerMessage = 3;
             tokensPerName = 1;
-        } else if (modelName.startsWith("gpt-3.5-turbo")) {
-            tokensPerMessage = 4; // every message follows <|start|>{role/name}\n{content}<|end|>\n
-            tokensPerName = -1; // if there's a name, the role is omitted
         }
+
         int sum = 0;
-        for (Message message : messages) {
+        for (final Message message : messages) {
             sum += tokensPerMessage;
             sum += encoding.countTokens(message.getContent());
             sum += encoding.countTokens(message.getRole());
-            if (StrUtil.isNotBlank(message.getName())) {
+            if (!StringUtils.isEmpty(message.getName())) {
                 sum += encoding.countTokens(message.getName());
                 sum += tokensPerName;
             }
         }
+
+        // every reply is primed with <|start|>assistant<|message|>
         sum += 3;
+
         return sum;
     }
-}
+
+    /**
+     * 计算tokens
+     *
+     * @param modelName 模型名称
+     * @param messages  消息列表
+     * @return 计算出的tokens数量
+     */
+
+    public static int tokens(String modelName, List<Message> messages) {
+        return tokens(messages, modelName);
+    }
+}