feat: rerank modal select and weight (#4164)

packages/global/core/ai/model.d.ts

@@ -72,7 +72,7 @@ export type EmbeddingModelItemType = PriceType &
     queryConfig?: Record<string, any>; // Custom parameters for query
   };
 
-export type ReRankModelItemType = PriceType &
+export type RerankModelItemType = PriceType &
   BaseModelItemType & {
     type: ModelTypeEnum.rerank;
   };

packages/global/core/app/type.d.ts

@@ -71,6 +71,19 @@ export type AppDetailType = AppSchema & {
   permission: AppPermission;
 };
 
+export type AppDatasetSearchParamsType = {
+  searchMode: `${DatasetSearchModeEnum}`;
+  limit?: number; // limit max tokens
+  similarity?: number;
+
+  usingReRank?: boolean;
+  rerankModel?: string;
+  rerankWeight?: number;
+
+  datasetSearchUsingExtensionQuery?: boolean;
+  datasetSearchExtensionModel?: string;
+  datasetSearchExtensionBg?: string;
+};
 export type AppSimpleEditFormType = {
   // templateId: string;
   aiSettings: {
@@ -88,14 +101,7 @@ export type AppSimpleEditFormType = {
   };
   dataset: {
     datasets: SelectedDatasetType;
-    searchMode: `${DatasetSearchModeEnum}`;
-    similarity?: number;
-    limit?: number;
-    usingReRank?: boolean;
-    datasetSearchUsingExtensionQuery?: boolean;
-    datasetSearchExtensionModel?: string;
-    datasetSearchExtensionBg?: string;
-  };
+  } & AppDatasetSearchParamsType;
   selectedTools: FlowNodeTemplateType[];
   chatConfig: AppChatConfigType;
 };

packages/global/core/app/utils.ts

@@ -24,9 +24,11 @@ export const getDefaultAppForm = (): AppSimpleEditFormType => {
     dataset: {
       datasets: [],
       similarity: 0.4,
-      limit: 1500,
+      limit: 3000,
       searchMode: DatasetSearchModeEnum.embedding,
       usingReRank: false,
+      rerankModel: '',
+      rerankWeight: 0.5,
       datasetSearchUsingExtensionQuery: true,
       datasetSearchExtensionBg: ''
     },
@@ -106,10 +108,20 @@ export const appWorkflow2Form = ({
       defaultAppForm.dataset.searchMode =
         findInputValueByKey(node.inputs, NodeInputKeyEnum.datasetSearchMode) ||
         DatasetSearchModeEnum.embedding;
+      // Rerank
       defaultAppForm.dataset.usingReRank = !!findInputValueByKey(
         node.inputs,
         NodeInputKeyEnum.datasetSearchUsingReRank
       );
+      defaultAppForm.dataset.rerankModel = findInputValueByKey(
+        node.inputs,
+        NodeInputKeyEnum.datasetSearchRerankModel
+      );
+      defaultAppForm.dataset.rerankWeight = findInputValueByKey(
+        node.inputs,
+        NodeInputKeyEnum.datasetSearchRerankWeight
+      );
+      // Query extension
       defaultAppForm.dataset.datasetSearchUsingExtensionQuery = findInputValueByKey(
         node.inputs,
         NodeInputKeyEnum.datasetSearchUsingExtensionQuery

packages/global/core/workflow/constants.ts

@@ -155,6 +155,8 @@ export enum NodeInputKeyEnum {
   datasetMaxTokens = 'limit',
   datasetSearchMode = 'searchMode',
   datasetSearchUsingReRank = 'usingReRank',
+  datasetSearchRerankWeight = 'rerankWeight',
+  datasetSearchRerankModel = 'rerankModel',
   datasetSearchUsingExtensionQuery = 'datasetSearchUsingExtensionQuery',
   datasetSearchExtensionModel = 'datasetSearchExtensionModel',
   datasetSearchExtensionBg = 'datasetSearchExtensionBg',

packages/global/core/workflow/template/system/datasetSearch.ts

@@ -64,6 +64,7 @@ export const DatasetSearchModule: FlowNodeTemplateType = {
       valueType: WorkflowIOValueTypeEnum.string,
       value: DatasetSearchModeEnum.embedding
     },
+    // Rerank
     {
       key: NodeInputKeyEnum.datasetSearchUsingReRank,
       renderTypeList: [FlowNodeInputTypeEnum.hidden],
@@ -71,6 +72,20 @@ export const DatasetSearchModule: FlowNodeTemplateType = {
       valueType: WorkflowIOValueTypeEnum.boolean,
       value: false
     },
+    {
+      key: NodeInputKeyEnum.datasetSearchRerankModel,
+      renderTypeList: [FlowNodeInputTypeEnum.hidden],
+      label: '',
+      valueType: WorkflowIOValueTypeEnum.string
+    },
+    {
+      key: NodeInputKeyEnum.datasetSearchRerankWeight,
+      renderTypeList: [FlowNodeInputTypeEnum.hidden],
+      label: '',
+      valueType: WorkflowIOValueTypeEnum.number,
+      value: 0.5
+    },
+    // Query Extension
     {
       key: NodeInputKeyEnum.datasetSearchUsingExtensionQuery,
       renderTypeList: [FlowNodeInputTypeEnum.hidden],
@@ -91,6 +106,7 @@ export const DatasetSearchModule: FlowNodeTemplateType = {
       valueType: WorkflowIOValueTypeEnum.string,
       value: ''
     },
+
     {
       key: NodeInputKeyEnum.authTmbId,
       renderTypeList: [FlowNodeInputTypeEnum.hidden],