support depth search

LLM-Red-Team · Jan 14, 2025 · db50132 · db50132
1 parent 9f2e172
commit db50132
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/src/api/controllers/chat.ts b/src/api/controllers/chat.ts
@@ -80,6 +80,7 @@ async function removeConversation(convId: string, ticket: string) {
 async function createCompletion(
   model = MODEL_NAME,
   messages: any[],
+  searchType: string = '',
   ticket: string,
   refConvId = '',
   retryCount = 0
@@ -129,7 +130,8 @@ async function createCompletion(
         sessionType: "text_chat",
         parentMsgId,
         params: {
-          "fileUploadBatchId": util.uuid()
+          "fileUploadBatchId": util.uuid(),
+          "searchType": searchType,
         },
         contents: messagesPrepare(messages, refs, !!refConvId),
       })
@@ -173,6 +175,7 @@ async function createCompletion(
 async function createCompletionStream(
   model = MODEL_NAME,
   messages: any[],
+  searchType: string = '',
   ticket: string,
   refConvId = '',
   retryCount = 0
@@ -220,7 +223,8 @@ async function createCompletionStream(
         sessionType: "text_chat",
         parentMsgId,
         params: {
-          "fileUploadBatchId": util.uuid()
+          "fileUploadBatchId": util.uuid(),
+          "searchType": searchType,
         },
         contents: messagesPrepare(messages, refs, !!refConvId),
       })

diff --git a/src/api/routes/chat.ts b/src/api/routes/chat.ts
@@ -17,11 +17,12 @@ export default {
       const tokens = chat.tokenSplit(request.headers.authorization);
       // 随机挑选一个ticket
       const token = _.sample(tokens);
-      const { model, conversation_id: convId, messages, stream } = request.body;
+      const { model, conversation_id: convId, messages, search_type, stream } = request.body;
       if (stream) {
         const stream = await chat.createCompletionStream(
           model,
           messages,
+          search_type,
           token,
           convId
         );
@@ -32,6 +33,7 @@ export default {
         return await chat.createCompletion(
           model,
           messages,
+          search_type,
           token,
           convId
         );