feat: 支持流式输出

Author: crazyurus

package.json

@@ -1,6 +1,6 @@
 {
   "name": "aigc-detector",
-  "version": "1.0.3",
+  "version": "1.0.4",
   "description": "Detect if content is generated by AI",
   "keywords": [
     "aigc",

src/cli/commands/chat.ts

@@ -42,8 +42,6 @@ class ChatCommand extends BaseCommand {
 
   static flags = {};
 
-  private lastMessage = 'How can I help you today?';
-
   private messages = new ChatMessageHistory();
 
   async run(): Promise<void> {
@@ -54,16 +52,25 @@ class ChatCommand extends BaseCommand {
         apiKey: config.apiKey,
         platform: config.platform as unknown as Platform
       });
-      const userDisplay = this.getDisplayContent(PromptRole.USER);
+      const aiDisplay = this.getDisplayContent(PromptRole.AI);
+      let lastMessage = 'How can I help you today?';
+
+      process.stdout.write(aiDisplay + lastMessage + '\n');
 
       // eslint-disable-next-line no-constant-condition
       while (true) {
-        const aiMessage = await this.addMessage(PromptRole.AI, this.lastMessage);
-        const userMessage = await this.getUserMessage(aiMessage + `\n${userDisplay}`);
-        const answer = await detector.chat(userMessage, await this.messages.getMessages());
+        const userMessage = await this.getUserMessage();
+        const stream = detector.chat(userMessage, await this.messages.getMessages());
+
+        process.stdout.write(aiDisplay);
+        stream.pipe(process.stdout);
+
+        lastMessage = await stream.getData();
+
+        process.stdout.write('\n');
 
         await this.addMessage(PromptRole.USER, userMessage);
-        this.lastMessage = answer;
+        await this.addMessage(PromptRole.AI, lastMessage);
       }
     } else {
       this.showHelp();
@@ -84,9 +91,11 @@ class ChatCommand extends BaseCommand {
     return chalk[roleDisplay.color](`[${roleDisplay.name}] `);
   }
 
-  private getUserMessage(aiMessage: string): Promise<string> {
+  private getUserMessage(): Promise<string> {
+    const userDisplay = this.getDisplayContent(PromptRole.USER);
+
     return new Promise<string>((resolve) => {
-      reader.question(aiMessage, resolve);
+      reader.question(userDisplay, resolve);
     });
   }
 }

src/core/index.ts

@@ -1,5 +1,7 @@
 import type { BaseMessage } from '@langchain/core/messages';
 
+import type Stream from './stream';
+
 import { PROMPT } from '../const';
 import { getPlatform, type Platform } from '../platform';
 import { getEnvConfig } from './env';
@@ -22,15 +24,14 @@ export class AIGC {
     this.platform = (env.platform as unknown as Platform) || options.platform;
   }
 
-  public async chat(content: string, messages: BaseMessage[]) {
+  public chat(content: string, messages: BaseMessage[]): Stream {
     const platform = getPlatform(this.platform);
-    const result = await platform.invoke(
+
+    return platform.stream(
       'You are a helpful assistant. Answer all questions to the best of your ability.',
       { content, messages },
       this.apiKey
     );
-
-    return result;
   }
 
   public async detect(content: string): Promise<ReturnType<typeof getDetectResult>> {

src/core/stream.ts

@@ -0,0 +1,23 @@
+import { Transform, type TransformCallback } from 'node:stream';
+
+class Stream extends Transform {
+  private data = '';
+
+  getData(): Promise<string> {
+    return new Promise((resolve) => {
+      this.on('close', () => {
+        resolve(this.data);
+      });
+    });
+  }
+
+  _transform(chunk: Buffer, encoding: BufferEncoding, callback: TransformCallback): void {
+    const data = chunk.toString();
+
+    this.data += data;
+
+    callback(null, data);
+  }
+}
+
+export default Stream;

src/platform/base.ts

@@ -4,19 +4,22 @@ import { ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemp
 import { ChatOpenAI } from '@langchain/openai';
 import { LLMChain } from 'langchain/chains';
 
+import Stream from '../core/stream';
+
 type InvokeParameter = Parameters<InstanceType<typeof LLMChain>['invoke']>[0];
 
 abstract class Platform {
   protected temperature = 0.7;
 
-  protected getChatModel(apiKey?: string): BaseLanguageModel {
+  protected getChatModel(apiKey?: string, streaming = false): BaseLanguageModel {
     return new ChatOpenAI({
       apiKey,
       configuration: {
         baseURL: `https://${this.server}/v1`
       },
       frequencyPenalty: 1,
       model: this.model,
+      streaming,
       temperature: this.temperature
     });
   }
@@ -39,6 +42,31 @@ abstract class Platform {
     return result.text;
   }
 
+  public stream(prompt: string, params: InvokeParameter, apiKey?: string): Stream {
+    const promptTemplate = this.getPrompt(prompt);
+    const chain = new LLMChain({
+      llm: this.getChatModel(apiKey, true),
+      prompt: promptTemplate
+    });
+    const stream = new Stream();
+
+    chain
+      .invoke(params, {
+        callbacks: [
+          {
+            handleLLMNewToken(token: string) {
+              stream.write(token);
+            }
+          }
+        ]
+      })
+      .then(() => {
+        stream.destroy();
+      });
+
+    return stream;
+  }
+
   protected abstract model: string;
 
   public abstract name: string;

src/platform/minimax.ts

@@ -11,11 +11,12 @@ class MiniMax extends Platform {
 
   protected server = 'api.minimax.chat';
 
-  protected getChatModel(apiKey?: string): BaseLanguageModel {
+  protected getChatModel(apiKey?: string, streaming = false): BaseLanguageModel {
     return new ChatMinimax({
       minimaxApiKey: apiKey,
       minimaxGroupId: '1782658868262748274',
       model: this.model,
+      streaming,
       temperature: this.temperature
     });
   }

src/platform/tongyi.ts

@@ -11,10 +11,11 @@ class TongYi extends Platform {
 
   protected server = 'open.bigmodel.cn';
 
-  protected getChatModel(apiKey?: string): BaseLanguageModel {
+  protected getChatModel(apiKey?: string, streaming = false): BaseLanguageModel {
     return new ChatAlibabaTongyi({
       alibabaApiKey: apiKey,
       model: this.model,
+      streaming,
       temperature: this.temperature
     });
   }

src/platform/zhipu.ts

@@ -11,9 +11,10 @@ class ZhiPu extends Platform {
 
   protected server = 'open.bigmodel.cn';
 
-  protected getChatModel(apiKey?: string): BaseLanguageModel {
+  protected getChatModel(apiKey?: string, streaming?: boolean): BaseLanguageModel {
     return new ChatZhipuAI({
       model: this.model,
+      streaming,
       temperature: this.temperature,
       zhipuAIApiKey: apiKey
     });