✨ Add Stream Support

bytemate · Feb 14, 2023 · 555b1af · 555b1af
1 parent 14d0bed
commit 555b1af
Showing 1 changed file with 104 additions and 20 deletions.
diff --git a/src/index.ts b/src/index.ts
@@ -1,16 +1,18 @@
 import { PrismaClient } from "@prisma/client";
 // @ts-ignore
-import { ChatGPTAPIBrowser, SendMessageOptions } from "chatgpt";
+import { ChatGPTAPIBrowser, ChatResponse, SendMessageOptions } from "chatgpt";
 import { loadConfig, getAccessToken } from "./lib";
 import express from "express";
 import AsyncRetry from "async-retry";
 import { Queue } from "async-await-queue";
 import { randomUUID } from "crypto";
+import Keyv from "keyv";
 const prisma = new PrismaClient();
 // ChatGPT (not plus) is limited to 1 request one time.
 const mesasgeQueue = new Queue(1, 100);
 const config = loadConfig();
 const app = express();
+const kv = new Keyv();
 let chatGPTAPIBrowser: ChatGPTAPIBrowser;
 app.use(express.json());
 app.get(`/`, async (req, res) => {
@@ -20,23 +22,6 @@ app.get(`/`, async (req, res) => {
   });
 });
 
-app.post(`/message`, async (req, res) => {
-  try {
-    const { message } = req.body;
-    console.log(`Received message: ${message}`);
-    const reply = await sendMesasge(message);
-    return res.json({
-      response: reply.response,
-    });
-  } catch (e) {
-    console.error(e);
-    return res.status(500).json({
-      message: "Something went wrong",
-      error: `${e}`,
-    });
-  }
-});
-
 const getOrCreateConversationInfo = async (
   sessionId: string
 ): Promise<SendMessageOptions> => {
@@ -54,7 +39,12 @@ const getOrCreateConversationInfo = async (
     return {};
   }
 };
-const sendMesasge = async (message: string, sessionId?: string) => {
+const sendMesasge = async (
+  message: string,
+  sessionId?: string,
+  // if message id is provided, it will be used to store the partial response
+  mesasgeId?: string
+) => {
   let conversationInfo;
   if (sessionId) {
     conversationInfo = await getOrCreateConversationInfo(sessionId);
@@ -63,11 +53,53 @@ const sendMesasge = async (message: string, sessionId?: string) => {
   await mesasgeQueue.wait(jobId);
   const startTime = new Date().getTime();
   let response;
+  let endFlag = false;
+  conversationInfo = {
+    ...conversationInfo,
+    mesasgeId,
+  };
+  if (mesasgeId) {
+    console.log(`message set ${mesasgeId}`);
+    await kv.set(
+      mesasgeId,
+      {
+        response: "",
+        status: "process",
+      },
+      30 * 60 * 1000
+    );
+  }
   try {
-    response = await chatGPTAPIBrowser.sendMessage(message, conversationInfo);
+    response = await chatGPTAPIBrowser.sendMessage(message, {
+      ...conversationInfo,
+      messageId: mesasgeId,
+      onProgress: mesasgeId
+        ? async (partialResponse: ChatResponse) => {
+            await kv.set(
+              mesasgeId,
+              {
+                ...partialResponse,
+                status: endFlag ? "done" : "process",
+              },
+              30 * 60 * 1000
+            );
+          }
+        : undefined,
+    });
+    endFlag = true;
     console.log(response);
     console.log(`Response: ${response}`);
   } catch (e) {
+    if (mesasgeId) {
+      await kv.set(
+        mesasgeId,
+        {
+          response: "",
+          status: "error",
+        },
+        30 * 60 * 1000
+      );
+    }
     console.error(e);
     throw e;
   } finally {
@@ -103,11 +135,51 @@ const sendMesasge = async (message: string, sessionId?: string) => {
   });
   return response;
 };
+app.post(`/message`, async (req, res) => {
+  try {
+    const { message } = req.body;
+    const { stream } = req.headers;
+    console.log(`Received message: ${message}`);
+    if (stream == "enable") {
+      const messageId = randomUUID();
+      sendMesasge(message, undefined, messageId).catch((e) => {
+        console.error(e);
+        console.log(`Error while sending message ${messageId}`);
+      });
+      return res.json({
+        messageId,
+      });
+    }
+    const reply = await sendMesasge(message);
+    return res.json({
+      response: reply.response,
+    });
+  } catch (e) {
+    console.error(e);
+    return res.status(500).json({
+      message: "Something went wrong",
+      error: `${e}`,
+    });
+  }
+});
 app.post(`/message/:sessionId`, async (req, res) => {
   try {
     const { sessionId } = req.params;
     const { message } = req.body;
     console.log(`Received message: ${message} for session: ${sessionId}`);
+    const { stream } = req.headers;
+    if (stream == "enable") {
+      const messageId = randomUUID();
+      sendMesasge(message, sessionId, messageId).catch((e) => {
+        console.error(e);
+        console.log(
+          `Error while sending message ${messageId}, sessionId: ${sessionId}`
+        );
+      });
+      return res.json({
+        messageId,
+      });
+    }
     const response = await sendMesasge(message, sessionId);
     return res.json({
       response: response.response,
@@ -120,6 +192,18 @@ app.post(`/message/:sessionId`, async (req, res) => {
     });
   }
 });
+app.get("/message/:messageId", async (req, res) => {
+  const { messageId } = req.params;
+  console.log(`${messageId}`);
+  const response = await kv.get(messageId);
+  if (response) {
+    return res.json(response);
+  } else {
+    return res.status(404).json({
+      message: "Not found",
+    });
+  }
+});
 app.delete(`/message/:sessionId`, async (req, res) => {
   try {
     const { sessionId } = req.params;