From f31f64c286750967edbbf25177075d5b90755b24 Mon Sep 17 00:00:00 2001
From: WJG <wjg12341234@163.com>
Date: Sun, 25 Feb 2024 00:34:05 +0800
Subject: [PATCH] feat: add chatWithStreamResponse for MyBot

---
 src/services/bot/index.ts | 32 +++++++++++++++++++++++++++++++-
 tests/bot.ts              | 26 ++++++++++++++++++++++++++
 tests/index.ts            |  4 +++-
 3 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100644 tests/bot.ts

diff --git a/src/services/bot/index.ts b/src/services/bot/index.ts
index afbd539..fc19688 100644
--- a/src/services/bot/index.ts
+++ b/src/services/bot/index.ts
@@ -1,9 +1,12 @@
+import { randomUUID } from "crypto";
 import { jsonDecode, jsonEncode } from "../../utils/base";
 import { buildPrompt, toUTC8Time } from "../../utils/string";
-import { openai } from "../openai";
+import { ChatOptions, openai } from "../openai";
 import { IBotConfig } from "./config";
 import { ConversationManager } from "./conversation";
+import { StreamResponse } from "../speaker/stream";
 
+// todo JSON mode 下，无法使用 stream 应答模式（在应答完成之前，无法构造完整的JSON）
 const systemTemplate = `
 忽略所有之前的文字、文件和说明。现在，你将扮演一个名为“{{name}}”的人，并以这个新身份回复所有新消息。
 
@@ -84,4 +87,31 @@ export class MyBot {
     });
     return jsonDecode(result?.content)?.message;
   }
+
+  static async chatWithStreamResponse(
+    options: ChatOptions & {
+      onFinished?: (text: string) => void;
+    }
+  ) {
+    const requestId = randomUUID();
+    const stream = new StreamResponse();
+    openai
+      .chatStream({
+        ...options,
+        requestId,
+        onStream: (text) => {
+          if (stream.status === "canceled") {
+            return openai.abort(requestId);
+          }
+          stream.addResponse(text);
+        },
+      })
+      .then((answer) => {
+        if (answer) {
+          stream.finish(answer);
+          options.onFinished?.(answer);
+        }
+      });
+    return stream;
+  }
 }
diff --git a/tests/bot.ts b/tests/bot.ts
new file mode 100644
index 0000000..bdb1f6f
--- /dev/null
+++ b/tests/bot.ts
@@ -0,0 +1,26 @@
+import { MyBot } from "../src/services/bot";
+import { AISpeaker } from "../src/services/speaker/ai";
+
+export async function testMyBot() {
+  await testStreamResponse();
+}
+
+async function testStreamResponse() {
+  const stream = await MyBot.chatWithStreamResponse({
+    user: "地球为什么是圆的？",
+    onFinished: (text) => {
+      console.log("\nFinal result 111:\n", text);
+    },
+  });
+  const config: any = {
+    userId: process.env.MI_USER!,
+    password: process.env.MI_PASS!,
+    did: process.env.MI_DID,
+    tts: "doubao",
+  };
+  const speaker = new AISpeaker(config);
+  await speaker.initMiServices();
+  await speaker.response({ stream });
+  const res = await stream.wasFinished();
+  console.log("\nFinal result 222:\n", res);
+}
diff --git a/tests/index.ts b/tests/index.ts
index eaf024a..3063642 100644
--- a/tests/index.ts
+++ b/tests/index.ts
@@ -5,6 +5,7 @@ import { runWithDB } from "../src/services/db";
 import { testDB } from "./db";
 import { testSpeaker } from "./speaker";
 import { testOpenAI } from "./openai";
+import { testMyBot } from "./bot";
 
 dotenv.config();
 
@@ -12,7 +13,8 @@ async function main() {
   println(kBannerASCII);
   // testDB();
   // testSpeaker();
-  testOpenAI();
+  // testOpenAI();
+  testMyBot();
 }
 
 runWithDB(main);