feat(server): use native tokenizer impl (#6960)

### Benchmark

`yarn workspace @affine/server-native bench`

```
┌─────────┬────────────┬─────────┬────────────────────┬──────────┬─────────┐
│ (index) │ Task Name  │ ops/sec │ Average Time (ns)  │ Margin   │ Samples │
├─────────┼────────────┼─────────┼────────────────────┼──────────┼─────────┤
│ 0       │ 'tiktoken' │ '5'     │ 176932518.76000002 │ '±4.71%' │ 100     │
│ 1       │ 'native'   │ '16'    │ 61041597.51000003  │ '±0.60%' │ 100     │
└─────────┴────────────┴─────────┴────────────────────┴──────────┴─────────┘
```

packages/backend/server/src/fundamentals/storage/index.ts

@@ -18,7 +18,7 @@ registerStorageProvider('fs', (config, bucket) => {
 })
 export class StorageProviderModule {}
 
-export * from './native';
+export * from '../../native';
 export type {
   BlobInputType,
   BlobOutputType,

packages/backend/server/src/fundamentals/storage/providers/utils.ts

@@ -3,7 +3,7 @@ import { Readable } from 'node:stream';
 import { crc32 } from '@node-rs/crc32';
 import { getStreamAsBuffer } from 'get-stream';
 
-import { getMime } from '../native';
+import { getMime } from '../../../native';
 import { BlobInputType, PutObjectMetadata } from './provider';
 
 export async function toBuffer(input: BlobInputType): Promise<Buffer> {

packages/backend/server/src/native.ts

@@ -7,10 +7,10 @@ try {
   const require = createRequire(import.meta.url);
   serverNativeModule =
     process.arch === 'arm64'
-      ? require('../../../server-native.arm64.node')
+      ? require('../server-native.arm64.node')
       : process.arch === 'arm'
-        ? require('../../../server-native.armv7.node')
-        : require('../../../server-native.node');
+        ? require('../server-native.armv7.node')
+        : require('../server-native.node');
 }
 
 export const mergeUpdatesInApplyWay = serverNativeModule.mergeUpdatesInApplyWay;
@@ -30,3 +30,5 @@ export const mintChallengeResponse = async (resource: string, bits: number) => {
 };
 
 export const getMime = serverNativeModule.getMime;
+export const Tokenizer = serverNativeModule.Tokenizer;
+export const fromModelName = serverNativeModule.fromModelName;

packages/backend/server/src/plugins/copilot/prompt.ts

@@ -1,7 +1,7 @@
+import { type Tokenizer } from '@affine/server-native';
 import { Injectable, Logger } from '@nestjs/common';
 import { AiPrompt, PrismaClient } from '@prisma/client';
 import Mustache from 'mustache';
-import { Tiktoken } from 'tiktoken';
 
 import {
   getTokenEncoder,
@@ -27,7 +27,7 @@ function extractMustacheParams(template: string) {
 
 export class ChatPrompt {
   private readonly logger = new Logger(ChatPrompt.name);
-  public readonly encoder?: Tiktoken;
+  public readonly encoder: Tokenizer | null;
   private readonly promptTokenSize: number;
   private readonly templateParamKeys: string[] = [];
   private readonly templateParams: PromptParams = {};
@@ -53,8 +53,7 @@ export class ChatPrompt {
   ) {
     this.encoder = getTokenEncoder(model);
     this.promptTokenSize =
-      this.encoder?.encode_ordinary(messages.map(m => m.content).join('') || '')
-        .length || 0;
+      this.encoder?.count(messages.map(m => m.content).join('') || '') || 0;
     this.templateParamKeys = extractMustacheParams(
       messages.map(m => m.content).join('')
     );
@@ -86,7 +85,7 @@ export class ChatPrompt {
   }
 
   encode(message: string) {
-    return this.encoder?.encode_ordinary(message).length || 0;
+    return this.encoder?.count(message) || 0;
   }
 
   private checkParams(params: PromptParams, sessionId?: string) {
@@ -129,10 +128,6 @@ export class ChatPrompt {
       content: Mustache.render(content, params),
     }));
   }
-
-  free() {
-    this.encoder?.free();
-  }
 }
 
 @Injectable()

packages/backend/server/src/plugins/copilot/session.ts

@@ -164,7 +164,6 @@ export class ChatSession implements AsyncDisposable {
   }
 
   async [Symbol.asyncDispose]() {
-    this.state.prompt.free();
     await this.save?.();
   }
 }
@@ -323,7 +322,7 @@ export class ChatSessionService {
   ): number {
     const encoder = getTokenEncoder(model);
     return messages
-      .map(m => encoder?.encode_ordinary(m.content).length || 0)
+      .map(m => encoder?.count(m.content) ?? 0)
       .reduce((total, length) => total + length, 0);
   }
 

packages/backend/server/src/plugins/copilot/types.ts

@@ -1,13 +1,9 @@
+import { type Tokenizer } from '@affine/server-native';
 import { AiPromptRole } from '@prisma/client';
 import type { ClientOptions as OpenAIClientOptions } from 'openai';
-import {
-  encoding_for_model,
-  get_encoding,
-  Tiktoken,
-  TiktokenModel,
-} from 'tiktoken';
 import { z } from 'zod';
 
+import { fromModelName } from '../../native';
 import type { ChatPrompt } from './prompt';
 import type { FalConfig } from './providers/fal';
 
@@ -37,17 +33,17 @@ export enum AvailableModels {
 
 export type AvailableModel = keyof typeof AvailableModels;
 
-export function getTokenEncoder(model?: string | null): Tiktoken | undefined {
-  if (!model) return undefined;
+export function getTokenEncoder(model?: string | null): Tokenizer | null {
+  if (!model) return null;
   const modelStr = AvailableModels[model as AvailableModel];
-  if (!modelStr) return undefined;
+  if (!modelStr) return null;
   if (modelStr.startsWith('gpt')) {
-    return encoding_for_model(modelStr as TiktokenModel);
+    return fromModelName(modelStr);
   } else if (modelStr.startsWith('dall')) {
     // dalle don't need to calc the token
-    return undefined;
+    return null;
   } else {
-    return get_encoding('cl100k_base');
+    return fromModelName('gpt-4-turbo-preview');
   }
 }