feat: chat quote

src/pages/api/openapi/chat/chat.ts

@@ -6,15 +6,19 @@ import { ChatItemSimpleType } from '@/types/chat';
 import { jsonRes } from '@/service/response';
 import { ChatModelMap, ModelVectorSearchModeMap } from '@/constants/model';
 import { pushChatBill } from '@/service/events/pushBill';
-import { searchKb } from '@/service/plugins/searchKb';
 import { ChatRoleEnum } from '@/constants/chat';
 import { withNextCors } from '@/service/utils/tools';
 import { BillTypeEnum } from '@/constants/user';
 import { sensitiveCheck } from '@/service/api/text';
+import { NEW_CHATID_HEADER } from '@/constants/chat';
+import { Types } from 'mongoose';
+import { appKbSearch } from '../kb/appKbSearch';
 
 /* 发送提示词 */
 export default withNextCors(async function handler(req: NextApiRequest, res: NextApiResponse) {
-  let step = 0; // step=1时，表示开始了流响应
+  res.on('close', () => {
+    res.end();
+  });
   res.on('error', () => {
     console.log('error: ', 'request error');
     res.end();
@@ -70,7 +74,7 @@ export default withNextCors(async function handler(req: NextApiRequest, res: Nex
 
     // 使用了知识库搜索
     if (model.chat.relatedKbs.length > 0) {
-      const { code, searchPrompts } = await searchKb({
+      const { code, searchPrompts } = await appKbSearch({
         prompts,
         similarity: ModelVectorSearchModeMap[model.chat.searchMode]?.similarity,
         model,
@@ -109,6 +113,10 @@ export default withNextCors(async function handler(req: NextApiRequest, res: Nex
       2
     );
 
+    // get conversationId. create a newId if it is null
+    const conversationId = chatId || String(new Types.ObjectId());
+    !chatId && res?.setHeader(NEW_CHATID_HEADER, conversationId);
+
     // 发出请求
     const { streamResponse, responseMessages, responseText, totalTokens } =
       await modelServiceToolMap[model.chat.chatModel].chatCompletion({
@@ -117,30 +125,41 @@ export default withNextCors(async function handler(req: NextApiRequest, res: Nex
         messages: prompts,
         stream: isStream,
         res,
-        chatId
+        chatId: conversationId
       });
 
     console.log('api response time:', `${(Date.now() - startTime) / 1000}s`);
 
-    let textLen = 0;
-    let tokens = totalTokens;
+    if (res.closed) return res.end();
 
-    if (isStream) {
-      step = 1;
-      const { finishMessages, totalTokens } = await resStreamResponse({
-        model: model.chat.chatModel,
-        res,
-        chatResponse: streamResponse,
-        prompts
-      });
-      textLen = finishMessages.map((item) => item.value).join('').length;
-      tokens = totalTokens;
-    } else {
-      textLen = responseMessages.map((item) => item.value).join('').length;
-      jsonRes(res, {
-        data: responseText
-      });
-    }
+    const { textLen = 0, tokens = totalTokens } = await (async () => {
+      if (isStream) {
+        try {
+          const { finishMessages, totalTokens } = await resStreamResponse({
+            model: model.chat.chatModel,
+            res,
+            chatResponse: streamResponse,
+            prompts
+          });
+          res.end();
+          return {
+            textLen: finishMessages.map((item) => item.value).join('').length,
+            tokens: totalTokens
+          };
+        } catch (error) {
+          res.end();
+          console.log('error，结束', error);
+        }
+      } else {
+        jsonRes(res, {
+          data: responseText
+        });
+        return {
+          textLen: responseMessages.map((item) => item.value).join('').length
+        };
+      }
+      return {};
+    })();
 
     pushChatBill({
       isPay: true,
@@ -151,16 +170,10 @@ export default withNextCors(async function handler(req: NextApiRequest, res: Nex
       type: BillTypeEnum.openapiChat
     });
   } catch (err: any) {
-    if (step === 1) {
-      // 直接结束流
-      res.end();
-      console.log('error，结束');
-    } else {
-      res.status(500);
-      jsonRes(res, {
-        code: 500,
-        error: err
-      });
-    }
+    res.status(500);
+    jsonRes(res, {
+      code: 500,
+      error: err
+    });
   }
 });