feat(core): enhance getEncoding function with global proxy support an…

…d timeout management
ChatLunaLab · Dec 30, 2024 · cf0691e · cf0691e
1 parent 8cff498
commit cf0691e
Showing 1 changed file with 28 additions and 23 deletions.
diff --git a/packages/core/src/llm-core/utils/tiktoken.ts b/packages/core/src/llm-core/utils/tiktoken.ts
@@ -5,7 +5,10 @@ import {
     TiktokenEncoding,
     TiktokenModel
 } from 'js-tiktoken/lite'
-import { chatLunaFetch } from 'koishi-plugin-chatluna/utils/request'
+import {
+    chatLunaFetch,
+    globalProxyAddress
+} from 'koishi-plugin-chatluna/utils/request'
 
 const cache: Record<string, TiktokenBPE> = {}
 
@@ -16,20 +19,38 @@ export async function getEncoding(
         extendedSpecialTokens?: Record<string, number>
     }
 ) {
+    options = options ?? {}
+
+    let timeout: NodeJS.Timeout
+
+    if (options.signal == null) {
+        const abortController = new AbortController()
+
+        options.signal = abortController.signal
+
+        timeout = setTimeout(() => abortController.abort(), 1000 * 5)
+    }
+
     if (!(encoding in cache)) {
-        cache[encoding] = await chatLunaFetch(
-            `https://tiktoken.pages.dev/js/${encoding}.json`,
-            {
-                signal: options?.signal
-            }
-        )
+        const url =
+            globalProxyAddress.length > 0
+                ? `https://tiktoken.pages.dev/js/${encoding}.json`
+                : `https://jsd.onmicrosoft.cn/npm/tiktoken@latest/encoders/${encoding}.json`
+
+        cache[encoding] = await chatLunaFetch(url, {
+            signal: options?.signal
+        })
             .then((res) => res.json() as unknown as TiktokenBPE)
             .catch((e) => {
                 delete cache[encoding]
                 throw e
             })
     }
 
+    if (timeout != null) {
+        clearTimeout(timeout)
+    }
+
     return new Tiktoken(cache[encoding], options?.extendedSpecialTokens)
 }
 
@@ -40,23 +61,7 @@ export async function encodingForModel(
         extendedSpecialTokens?: Record<string, number>
     }
 ) {
-    options = options ?? {}
-
-    let timeout: NodeJS.Timeout
-
-    if (options.signal == null) {
-        const abortController = new AbortController()
-
-        options.signal = abortController.signal
-
-        timeout = setTimeout(() => abortController.abort(), 1000 * 5)
-    }
-
     const result = await getEncoding(getEncodingNameForModel(model), options)
 
-    if (timeout != null) {
-        clearTimeout(timeout)
-    }
-
     return result
 }