From cf0691eb952aadc79089f8dd50b91bd3f954a53f Mon Sep 17 00:00:00 2001
From: dingyi <dingyi222666@foxmail.com>
Date: Mon, 30 Dec 2024 09:57:41 +0800
Subject: [PATCH] feat(core): enhance getEncoding function with global proxy
 support and timeout management

---
 packages/core/src/llm-core/utils/tiktoken.ts | 51 +++++++++++---------
 1 file changed, 28 insertions(+), 23 deletions(-)

diff --git a/packages/core/src/llm-core/utils/tiktoken.ts b/packages/core/src/llm-core/utils/tiktoken.ts
index e238663f..fb877f57 100644
--- 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,13 +19,27 @@ 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]
@@ -30,6 +47,10 @@ export async function getEncoding(
             })
     }
 
+    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
 }