feat: rerank modal select and weight (#4164)

This commit is contained in:
Archer
2025-03-14 14:49:27 +08:00
committed by GitHub
parent 561a496f80
commit d8712d4092
36 changed files with 282 additions and 178 deletions

View File

@@ -6,7 +6,7 @@ import type {
EmbeddingModelItemType,
AudioSpeechModels,
STTModelType,
ReRankModelItemType
RerankModelItemType
} from '../../../core/ai/model.d';
import { SubTypeEnum } from '../../../support/wallet/sub/constants';
@@ -35,7 +35,7 @@ export type FastGPTConfigFileType = {
// Abandon
llmModels?: ChatModelItemType[];
vectorModels?: EmbeddingModelItemType[];
reRankModels?: ReRankModelItemType[];
reRankModels?: RerankModelItemType[];
audioSpeechModels?: TTSModelType[];
whisperModel?: STTModelType;
};

View File

@@ -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;
};

View File

@@ -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;
};

View File

@@ -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

View File

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

View File

@@ -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],