⚡️ perf: 优化长文本时的渲染性能 (lobehub#3754)

* ⚡️ perf: 优化渲染长文本时的渲染性能

* 👷 build: fix vercel build

* ⚡️ perf: 提高 smooth 每帧速度

* ✅ test: add test for tokenizer edge runtime

* 💚 build: fix build

* ♻️ refactor: refactor with webapi

* 🚨 ci: improve lint

* ✅ test: fix test

* ⚡️ perf: try o200k_base

src/app/(main)/profile/[[...slugs]]/Client.tsx

@@ -56,6 +56,7 @@ export const useStyles = createStyles(
         border-radius: unset;
       `,
     }) as Partial<{
+      // eslint-disable-next-line unused-imports/no-unused-vars
       [k in keyof ElementsConfig]: any;
     }>,
 );

src/app/webapi/tokenizer/index.test.ts

@@ -0,0 +1,32 @@
+// @vitest-environment edge-runtime
+import { describe, expect, it } from 'vitest';
+
+import { POST } from './route';
+
+describe('tokenizer Route', () => {
+  it('count hello world', async () => {
+    const txt = 'Hello, world!';
+    const request = new Request('https://test.com', {
+      method: 'POST',
+      body: txt,
+    });
+
+    const response = await POST(request);
+
+    const data = await response.json();
+    expect(data.count).toEqual(4);
+  });
+
+  it('count Chinese', async () => {
+    const txt = '今天天气真好';
+    const request = new Request('https://test.com', {
+      method: 'POST',
+      body: txt,
+    });
+
+    const response = await POST(request);
+
+    const data = await response.json();
+    expect(data.count).toEqual(5);
+  });
+});

src/app/webapi/tokenizer/route.ts

@@ -0,0 +1,8 @@
+import { encode } from 'gpt-tokenizer/encoding/o200k_base';
+import { NextResponse } from 'next/server';
+
+export const POST = async (req: Request) => {
+  const str = await req.text();
+
+  return NextResponse.json({ count: encode(str).length });
+};

src/database/server/models/session.ts

@@ -170,7 +170,7 @@ export class SessionModel {
 
     if (!result) return;
 
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars,unused-imports/no-unused-vars
     const { agent, clientId, ...session } = result;
     const sessionId = this.genId();
 

src/hooks/useTokenCount.ts

@@ -1,20 +1,32 @@
-import { startTransition, useEffect, useState } from 'react';
+import { debounce } from 'lodash-es';
+import { startTransition, useCallback, useEffect, useState } from 'react';
 
 import { encodeAsync } from '@/utils/tokenizer';
 
 export const useTokenCount = (input: string = '') => {
   const [value, setNum] = useState(0);
 
-  useEffect(() => {
-    startTransition(() => {
-      encodeAsync(input || '')
+  const debouncedEncode = useCallback(
+    debounce((text: string) => {
+      encodeAsync(text)
         .then(setNum)
         .catch(() => {
-          // 兜底采用字符数
-          setNum(input.length);
+          setNum(text.length);
         });
+    }, 300),
+    [],
+  );
+
+  useEffect(() => {
+    startTransition(() => {
+      debouncedEncode(input || '');
     });
-  }, [input]);
+
+    // 清理函数
+    return () => {
+      debouncedEncode.cancel();
+    };
+  }, [input, debouncedEncode]);
 
   return value;
 };

src/layout/AuthProvider/Clerk/useAppearance.ts

@@ -89,6 +89,7 @@ export const useStyles = createStyles(
         order: -1;
       `,
     }) as Partial<{
+      // eslint-disable-next-line unused-imports/no-unused-vars
       [k in keyof ElementsConfig]: any;
     }>,
 );

src/server/context.ts

@@ -62,7 +62,9 @@ export const createContext = async (request: NextRequest): Promise<Context> => {
         userId = session.user.id;
       }
       return createContextInner({ authorizationHeader: authorization, nextAuth: auth, userId });
-    } catch {}
+    } catch (e) {
+      console.error('next auth err', e);
+    }
   }
 
   return createContextInner({ authorizationHeader: authorization, userId });

src/server/routers/edge/index.ts

@@ -1,5 +1,5 @@
 /**
- * This file contains the root router of Lobe Chat tRPC-backend
+ * This file contains the edge router of Lobe Chat tRPC-backend
  */
 import { publicProcedure, router } from '@/libs/trpc';
 

src/services/user/client.ts

@@ -50,7 +50,7 @@ export class ClientService implements IUserService {
     await this.preferenceStorage.saveToLocalStorage(preference);
   }
 
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars,unused-imports/no-unused-vars
   async updateGuide(guide: Partial<UserGuide>) {
     throw new Error('Method not implemented.');
   }

src/types/worker.d.ts

@@ -0,0 +1,7 @@
+declare module '*.worker.ts' {
+  class WebpackWorker extends Worker {
+    constructor();
+  }
+
+  export default WebpackWorker;
+}

src/utils/fetch/__tests__/fetchSSE.test.ts

@@ -139,8 +139,8 @@ describe('fetchSSE', () => {
       onFinish: mockOnFinish,
     });
 
-    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(1, { text: 'He', type: 'text' });
-    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(2, { text: 'llo World', type: 'text' });
+    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(1, { text: 'Hell', type: 'text' });
+    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(2, { text: 'o World', type: 'text' });
     // more assertions for each character...
     expect(mockOnFinish).toHaveBeenCalledWith('Hello World', {
       observationId: null,
@@ -232,8 +232,8 @@ describe('fetchSSE', () => {
       signal: abortController.signal,
     });
 
-    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(1, { text: 'He', type: 'text' });
-    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(2, { text: 'llo World', type: 'text' });
+    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(1, { text: 'Hell', type: 'text' });
+    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(2, { text: 'o World', type: 'text' });
 
     expect(mockOnFinish).toHaveBeenCalledWith('Hello World', {
       type: 'done',

src/utils/fetch/fetchSSE.ts

@@ -44,6 +44,10 @@ export interface FetchSSEOptions {
   smoothing?: boolean;
 }
 
+const START_ANIMATION_SPEED = 4;
+
+const END_ANIMATION_SPEED = 15;
+
 const createSmoothMessage = (params: { onTextUpdate: (delta: string, text: string) => void }) => {
   let buffer = '';
   // why use queue: https://shareg.pt/GLBrjpK
@@ -64,7 +68,7 @@ const createSmoothMessage = (params: { onTextUpdate: (delta: string, text: strin
 
   // define startAnimation function to display the text in buffer smooth
   // when you need to start the animation, call this function
-  const startAnimation = (speed = 2) =>
+  const startAnimation = (speed = START_ANIMATION_SPEED) =>
     new Promise<void>((resolve) => {
       if (isAnimationActive) {
         resolve();
@@ -137,7 +141,7 @@ const createSmoothToolCalls = (params: {
     }
   };
 
-  const startAnimation = (index: number, speed = 2) =>
+  const startAnimation = (index: number, speed = START_ANIMATION_SPEED) =>
     new Promise<void>((resolve) => {
       if (isAnimationActives[index]) {
         resolve();
@@ -191,7 +195,7 @@ const createSmoothToolCalls = (params: {
     });
   };
 
-  const startAnimations = async (speed = 2) => {
+  const startAnimations = async (speed = START_ANIMATION_SPEED) => {
     const pools = toolCallsBuffer.map(async (_, index) => {
       if (outputQueues[index].length > 0 && !isAnimationActives[index]) {
         await startAnimation(index, speed);
@@ -365,11 +369,11 @@ export const fetchSSE = async (url: string, options: RequestInit & FetchSSEOptio
       const observationId = response.headers.get(LOBE_CHAT_OBSERVATION_ID);
 
       if (textController.isTokenRemain()) {
-        await textController.startAnimation(15);
+        await textController.startAnimation(END_ANIMATION_SPEED);
       }
 
       if (toolCallsController.isTokenRemain()) {
-        await toolCallsController.startAnimations(15);
+        await toolCallsController.startAnimations(END_ANIMATION_SPEED);
       }
 
       await options?.onFinish?.(output, { observationId, toolCalls, traceId, type: finishedType });

src/utils/tokenizer/client.ts

@@ -0,0 +1,35 @@
+let worker: Worker | null = null;
+
+const getWorker = () => {
+  if (!worker && typeof Worker !== 'undefined') {
+    worker = new Worker(new URL('tokenizer.worker.ts', import.meta.url));
+  }
+  return worker;
+};
+
+export const clientEncodeAsync = (str: string): Promise<number> =>
+  new Promise((resolve, reject) => {
+    const worker = getWorker();
+
+    if (!worker) {
+      // 如果 WebWorker 不可用，回退到字符串计算
+      resolve(str.length);
+      return;
+    }
+
+    const id = Date.now().toString();
+
+    const handleMessage = (event: MessageEvent) => {
+      if (event.data.id === id) {
+        worker.removeEventListener('message', handleMessage);
+        if (event.data.error) {
+          reject(new Error(event.data.error));
+        } else {
+          resolve(event.data.result);
+        }
+      }
+    };
+
+    worker.addEventListener('message', handleMessage);
+    worker.postMessage({ id, str });
+  });

src/utils/tokenizer/index.ts

@@ -0,0 +1,15 @@
+export const encodeAsync = async (str: string): Promise<number> => {
+  if (str.length === 0) return 0;
+
+  // 50_000 is the limit of the client
+  // if the string is longer than 100_000, we will use the server
+  if (str.length <= 50_000) {
+    const { clientEncodeAsync } = await import('./client');
+
+    return await clientEncodeAsync(str);
+  } else {
+    const { serverEncodeAsync } = await import('./server');
+
+    return await serverEncodeAsync(str);
+  }
+};

src/utils/tokenizer/server.ts

@@ -0,0 +1,11 @@
+export const serverEncodeAsync = async (str: string): Promise<number> => {
+  try {
+    const res = await fetch('/webapi/tokenizer', { body: str, method: 'POST' });
+    const data = await res.json();
+
+    return data.count;
+  } catch (e) {
+    console.error('serverEncodeAsync:', e);
+    return str.length;
+  }
+};

src/utils/tokenizer/tokenizer.worker.ts

@@ -0,0 +1,14 @@
+addEventListener('message', async (event) => {
+  const { id, str } = event.data;
+  try {
+    const { encode } = await import('gpt-tokenizer');
+
+    console.time('client tokenizer');
+    const tokenCount = encode(str).length;
+    console.timeEnd('client tokenizer');
+
+    postMessage({ id, result: tokenCount });
+  } catch (error) {
+    postMessage({ error: (error as Error).message, id });
+  }
+});