feat: chat quote

src/pages/api/chat/chat.ts

@@ -2,19 +2,24 @@ import type { NextApiRequest, NextApiResponse } from 'next';
 import { connectToDatabase } from '@/service/mongo';
 import { authChat } from '@/service/utils/auth';
 import { modelServiceToolMap } from '@/service/utils/chat';
-import { ChatItemSimpleType } from '@/types/chat';
+import { ChatItemType } from '@/types/chat';
 import { jsonRes } from '@/service/response';
 import { ChatModelMap, ModelVectorSearchModeMap } from '@/constants/model';
 import { pushChatBill } from '@/service/events/pushBill';
 import { resStreamResponse } from '@/service/utils/chat';
-import { searchKb } from '@/service/plugins/searchKb';
+import { appKbSearch } from '../openapi/kb/appKbSearch';
 import { ChatRoleEnum } from '@/constants/chat';
 import { BillTypeEnum } from '@/constants/user';
 import { sensitiveCheck } from '@/service/api/text';
+import { NEW_CHATID_HEADER } from '@/constants/chat';
+import { saveChat } from './saveChat';
+import { Types } from 'mongoose';
 
 /* 发送提示词 */
 export default 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();
@@ -22,9 +27,9 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
 
   try {
     const { chatId, prompt, modelId } = req.body as {
-      prompt: ChatItemSimpleType;
+      prompt: [ChatItemType, ChatItemType];
       modelId: string;
-      chatId: '' | string;
+      chatId?: string;
     };
 
     if (!modelId || !prompt) {
@@ -44,42 +49,69 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
     const modelConstantsData = ChatModelMap[model.chat.chatModel];
 
     // 读取对话内容
-    const prompts = [...content, prompt];
-    let systemPrompts: {
-      obj: ChatRoleEnum;
-      value: string;
-    }[] = [];
+    const prompts = [...content, prompt[0]];
+    const {
+      code = 200,
+      systemPrompts = [],
+      quote = []
+    } = await (async () => {
+      // 使用了知识库搜索
+      if (model.chat.relatedKbs.length > 0) {
+        const { code, searchPrompts, rawSearch } = await appKbSearch({
+          model,
+          userId,
+          prompts,
+          similarity: ModelVectorSearchModeMap[model.chat.searchMode]?.similarity
+        });
 
-    // 使用了知识库搜索
-    if (model.chat.relatedKbs.length > 0) {
-      const { code, searchPrompts } = await searchKb({
-        userOpenAiKey,
-        prompts,
-        similarity: ModelVectorSearchModeMap[model.chat.searchMode]?.similarity,
-        model,
+        return {
+          code,
+          quote: rawSearch,
+          systemPrompts: searchPrompts
+        };
+      }
+      if (model.chat.systemPrompt) {
+        return {
+          systemPrompts: [
+            {
+              obj: ChatRoleEnum.System,
+              value: model.chat.systemPrompt
+            }
+          ]
+        };
+      }
+      return {};
+    })();
+
+    // get conversationId. create a newId if it is null
+    const conversationId = chatId || String(new Types.ObjectId());
+    !chatId && res?.setHeader(NEW_CHATID_HEADER, conversationId);
+
+    // search result is empty
+    if (code === 201) {
+      const response = systemPrompts[0]?.value;
+      await saveChat({
+        chatId,
+        newChatId: conversationId,
+        modelId,
+        prompts: [
+          prompt[0],
+          {
+            ...prompt[1],
+            quote: [],
+            value: response
+          }
+        ],
         userId
       });
-
-      // search result is empty
-      if (code === 201) {
-        return res.send(searchPrompts[0]?.value);
-      }
-
-      systemPrompts = searchPrompts;
-    } else if (model.chat.systemPrompt) {
-      systemPrompts = [
-        {
-          obj: ChatRoleEnum.System,
-          value: model.chat.systemPrompt
-        }
-      ];
+      return res.end(response);
     }
 
     prompts.splice(prompts.length - 3, 0, ...systemPrompts);
 
     // content check
     await sensitiveCheck({
-      input: [...systemPrompts, prompt].map((item) => item.value).join('')
+      input: [...systemPrompts, prompt[0]].map((item) => item.value).join('')
     });
 
     // 计算温度
@@ -87,54 +119,65 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
       2
     );
 
-    // 发出请求
+    // 发出 chat 请求
     const { streamResponse } = await modelServiceToolMap[model.chat.chatModel].chatCompletion({
       apiKey: userOpenAiKey || systemAuthKey,
       temperature: +temperature,
       messages: prompts,
       stream: true,
       res,
-      chatId
+      chatId: conversationId
     });
 
     console.log('api response time:', `${(Date.now() - startTime) / 1000}s`);
 
-    step = 1;
+    if (res.closed) return res.end();
 
-    const { totalTokens, finishMessages } = await resStreamResponse({
-      model: model.chat.chatModel,
-      res,
-      chatResponse: streamResponse,
-      prompts,
-      systemPrompt: showModelDetail
-        ? prompts
-            .filter((item) => item.obj === ChatRoleEnum.System)
-            .map((item) => item.value)
-            .join('\n')
-        : ''
-    });
-
-    // 只有使用平台的 key 才计费
-    pushChatBill({
-      isPay: !userOpenAiKey,
-      chatModel: model.chat.chatModel,
-      userId,
-      chatId,
-      textLen: finishMessages.map((item) => item.value).join('').length,
-      tokens: totalTokens,
-      type: BillTypeEnum.chat
-    });
-  } catch (err: any) {
-    if (step === 1) {
-      // 直接结束流
-      res.end();
-      console.log('error，结束');
-    } else {
-      res.status(500);
-      jsonRes(res, {
-        code: 500,
-        error: err
+    try {
+      const { totalTokens, finishMessages, responseContent } = await resStreamResponse({
+        model: model.chat.chatModel,
+        res,
+        chatResponse: streamResponse,
+        prompts
       });
+
+      // save chat
+      await saveChat({
+        chatId,
+        newChatId: conversationId,
+        modelId,
+        prompts: [
+          prompt[0],
+          {
+            ...prompt[1],
+            quote: showModelDetail ? quote : [],
+            value: responseContent
+          }
+        ],
+        userId
+      });
+
+      res.end();
+
+      // 只有使用平台的 key 才计费
+      pushChatBill({
+        isPay: !userOpenAiKey,
+        chatModel: model.chat.chatModel,
+        userId,
+        chatId: conversationId,
+        textLen: finishMessages.map((item) => item.value).join('').length,
+        tokens: totalTokens,
+        type: BillTypeEnum.chat
+      });
+    } catch (error) {
+      res.end();
+      console.log('error，结束', error);
     }
+  } catch (err: any) {
+    res.status(500);
+    jsonRes(res, {
+      code: 500,
+      error: err
+    });
   }
 }