Commit 257e795c authored by Joel's avatar Joel

feat: retrieval config

parent bafdc510
'use client'
import React from 'react'
import type { FC } from 'react'
import { useTranslation } from 'react-i18next'
import TopKItem from '@/app/components/base/param-item/top-k-item'
import ScoreThresholdItem from '@/app/components/base/param-item/score-threshold-item'
import RadioCard from '@/app/components/base/radio-card/simple'
import { RETRIEVE_TYPE } from '@/types/app'
import {
MultiPathRetrieval,
NTo1Retrieval,
} from '@/app/components/base/icons/src/public/common'
import type {
DatasetConfigs,
} from '@/models/debug'
import ModelSelector from '@/app/components/header/account-setting/model-provider-page/model-selector'
import { useModelListAndDefaultModelAndCurrentProviderAndModel } from '@/app/components/header/account-setting/model-provider-page/hooks'
type Props = {
datasetConfigs: DatasetConfigs
onChange: (configs: DatasetConfigs, isRetrievalModeChange?: boolean) => void
}
const ConfigContent: FC<Props> = ({
datasetConfigs,
onChange,
}) => {
const { t } = useTranslation()
const type = datasetConfigs.retrieval_model
const setType = (value: RETRIEVE_TYPE) => {
onChange({
...datasetConfigs,
retrieval_model: value,
}, true)
}
const {
modelList: rerankModelList,
defaultModel: rerankDefaultModel,
currentModel: isRerankDefaultModelVaild,
} = useModelListAndDefaultModelAndCurrentProviderAndModel(3)
const rerankModel = (() => {
if (datasetConfigs.reranking_model) {
return {
provider_name: datasetConfigs.reranking_model.reranking_provider_name,
model_name: datasetConfigs.reranking_model.reranking_model_name,
}
}
else if (rerankDefaultModel) {
return {
provider_name: rerankDefaultModel.provider.provider,
model_name: rerankDefaultModel.model,
}
}
})()
const handleParamChange = (key: string, value: number) => {
if (key === 'top_k') {
onChange({
...datasetConfigs,
top_k: value,
})
}
else if (key === 'score_threshold') {
onChange({
...datasetConfigs,
score_threshold: value,
})
}
}
const handleSwitch = (key: string, enable: boolean) => {
if (key === 'top_k')
return
onChange({
...datasetConfigs,
score_threshold_enabled: enable,
})
}
return (
<div>
<div className='mt-2 space-y-3'>
<RadioCard
icon={<NTo1Retrieval className='shrink-0 mr-3 w-9 h-9 rounded-lg' />}
title={t('appDebug.datasetConfig.retrieveOneWay.title')}
description={t('appDebug.datasetConfig.retrieveOneWay.description')}
isChosen={type === RETRIEVE_TYPE.oneWay}
onChosen={() => { setType(RETRIEVE_TYPE.oneWay) }}
/>
<RadioCard
icon={<MultiPathRetrieval className='shrink-0 mr-3 w-9 h-9 rounded-lg' />}
title={t('appDebug.datasetConfig.retrieveMultiWay.title')}
description={t('appDebug.datasetConfig.retrieveMultiWay.description')}
isChosen={type === RETRIEVE_TYPE.multiWay}
onChosen={() => { setType(RETRIEVE_TYPE.multiWay) }}
/>
</div>
{type === RETRIEVE_TYPE.multiWay && (
<>
<div className='mt-6'>
<div className='leading-[32px] text-[13px] font-medium text-gray-900'>{t('common.modelProvider.rerankModel.key')}</div>
<div>
<ModelSelector
defaultModel={rerankModel && { provider: rerankModel?.provider_name, model: rerankModel?.model_name }}
onSelect={(v) => {
onChange({
...datasetConfigs,
reranking_model: {
reranking_provider_name: v.provider,
reranking_model_name: v.model,
},
})
}}
modelList={rerankModelList}
/>
</div>
</div>
<div className='mt-4 space-y-4'>
<TopKItem
value={datasetConfigs.top_k}
onChange={handleParamChange}
enable={true}
/>
<ScoreThresholdItem
value={datasetConfigs.score_threshold}
onChange={handleParamChange}
enable={datasetConfigs.score_threshold_enabled}
hasSwitch={true}
onSwitchChange={handleSwitch}
/>
</div>
</>
)}
</div>
)
}
export default React.memo(ConfigContent)
...@@ -4,21 +4,14 @@ import { memo, useState } from 'react' ...@@ -4,21 +4,14 @@ import { memo, useState } from 'react'
import { useTranslation } from 'react-i18next' import { useTranslation } from 'react-i18next'
import { useContext } from 'use-context-selector' import { useContext } from 'use-context-selector'
import cn from 'classnames' import cn from 'classnames'
import ConfigContent from './config-content'
import { Settings04 } from '@/app/components/base/icons/src/vender/line/general' import { Settings04 } from '@/app/components/base/icons/src/vender/line/general'
import ConfigContext from '@/context/debug-configuration' import ConfigContext from '@/context/debug-configuration'
import TopKItem from '@/app/components/base/param-item/top-k-item'
import ScoreThresholdItem from '@/app/components/base/param-item/score-threshold-item'
import Modal from '@/app/components/base/modal' import Modal from '@/app/components/base/modal'
import Button from '@/app/components/base/button' import Button from '@/app/components/base/button'
import RadioCard from '@/app/components/base/radio-card/simple'
import { RETRIEVE_TYPE } from '@/types/app' import { RETRIEVE_TYPE } from '@/types/app'
import Toast from '@/app/components/base/toast' import Toast from '@/app/components/base/toast'
import { DATASET_DEFAULT } from '@/config' import { DATASET_DEFAULT } from '@/config'
import {
MultiPathRetrieval,
NTo1Retrieval,
} from '@/app/components/base/icons/src/public/common'
import ModelSelector from '@/app/components/header/account-setting/model-provider-page/model-selector'
import { useModelListAndDefaultModelAndCurrentProviderAndModel } from '@/app/components/header/account-setting/model-provider-page/hooks' import { useModelListAndDefaultModelAndCurrentProviderAndModel } from '@/app/components/header/account-setting/model-provider-page/hooks'
const ParamsConfig: FC = () => { const ParamsConfig: FC = () => {
...@@ -30,58 +23,11 @@ const ParamsConfig: FC = () => { ...@@ -30,58 +23,11 @@ const ParamsConfig: FC = () => {
} = useContext(ConfigContext) } = useContext(ConfigContext)
const [tempDataSetConfigs, setTempDataSetConfigs] = useState(datasetConfigs) const [tempDataSetConfigs, setTempDataSetConfigs] = useState(datasetConfigs)
const type = tempDataSetConfigs.retrieval_model
const setType = (value: RETRIEVE_TYPE) => {
setTempDataSetConfigs({
...tempDataSetConfigs,
retrieval_model: value,
})
}
const { const {
modelList: rerankModelList,
defaultModel: rerankDefaultModel, defaultModel: rerankDefaultModel,
currentModel: isRerankDefaultModelVaild, currentModel: isRerankDefaultModelVaild,
} = useModelListAndDefaultModelAndCurrentProviderAndModel(3) } = useModelListAndDefaultModelAndCurrentProviderAndModel(3)
const rerankModel = (() => {
if (tempDataSetConfigs.reranking_model) {
return {
provider_name: tempDataSetConfigs.reranking_model.reranking_provider_name,
model_name: tempDataSetConfigs.reranking_model.reranking_model_name,
}
}
else if (rerankDefaultModel) {
return {
provider_name: rerankDefaultModel.provider.provider,
model_name: rerankDefaultModel.model,
}
}
})()
const handleParamChange = (key: string, value: number) => {
if (key === 'top_k') {
setTempDataSetConfigs({
...tempDataSetConfigs,
top_k: value,
})
}
else if (key === 'score_threshold') {
setTempDataSetConfigs({
...tempDataSetConfigs,
score_threshold: value,
})
}
}
const handleSwitch = (key: string, enable: boolean) => {
if (key === 'top_k')
return
setTempDataSetConfigs({
...tempDataSetConfigs,
score_threshold_enabled: enable,
})
}
const isValid = () => { const isValid = () => {
let errMsg = '' let errMsg = ''
if (tempDataSetConfigs.retrieval_model === RETRIEVE_TYPE.multiWay) { if (tempDataSetConfigs.retrieval_model === RETRIEVE_TYPE.multiWay) {
...@@ -140,58 +86,11 @@ const ParamsConfig: FC = () => { ...@@ -140,58 +86,11 @@ const ParamsConfig: FC = () => {
wrapperClassName='z-50' wrapperClassName='z-50'
title={t('appDebug.datasetConfig.settingTitle')} title={t('appDebug.datasetConfig.settingTitle')}
> >
<div className='mt-2 space-y-3'> <ConfigContent
<RadioCard datasetConfigs={tempDataSetConfigs}
icon={<NTo1Retrieval className='shrink-0 mr-3 w-9 h-9 rounded-lg' />} onChange={setTempDataSetConfigs}
title={t('appDebug.datasetConfig.retrieveOneWay.title')} />
description={t('appDebug.datasetConfig.retrieveOneWay.description')}
isChosen={type === RETRIEVE_TYPE.oneWay}
onChosen={() => { setType(RETRIEVE_TYPE.oneWay) }}
/>
<RadioCard
icon={<MultiPathRetrieval className='shrink-0 mr-3 w-9 h-9 rounded-lg' />}
title={t('appDebug.datasetConfig.retrieveMultiWay.title')}
description={t('appDebug.datasetConfig.retrieveMultiWay.description')}
isChosen={type === RETRIEVE_TYPE.multiWay}
onChosen={() => { setType(RETRIEVE_TYPE.multiWay) }}
/>
</div>
{type === RETRIEVE_TYPE.multiWay && (
<>
<div className='mt-6'>
<div className='leading-[32px] text-[13px] font-medium text-gray-900'>{t('common.modelProvider.rerankModel.key')}</div>
<div>
<ModelSelector
defaultModel={rerankModel && { provider: rerankModel?.provider_name, model: rerankModel?.model_name }}
onSelect={(v) => {
setTempDataSetConfigs({
...tempDataSetConfigs,
reranking_model: {
reranking_provider_name: v.provider,
reranking_model_name: v.model,
},
})
}}
modelList={rerankModelList}
/>
</div>
</div>
<div className='mt-4 space-y-4'>
<TopKItem
value={tempDataSetConfigs.top_k}
onChange={handleParamChange}
enable={true}
/>
<ScoreThresholdItem
value={tempDataSetConfigs.score_threshold}
onChange={handleParamChange}
enable={tempDataSetConfigs.score_threshold_enabled}
hasSwitch={true}
onSwitchChange={handleSwitch}
/>
</div>
</>
)}
<div className='mt-6 flex justify-end'> <div className='mt-6 flex justify-end'>
<Button className='mr-2 flex-shrink-0' onClick={() => { <Button className='mr-2 flex-shrink-0' onClick={() => {
setOpen(false) setOpen(false)
......
'use client'
import type { FC } from 'react'
import React, { useCallback, useState } from 'react'
import { useTranslation } from 'react-i18next'
import cn from 'classnames'
import type { MultipleRetrievalConfig } from '../types'
import {
PortalToFollowElem,
PortalToFollowElemContent,
PortalToFollowElemTrigger,
} from '@/app/components/base/portal-to-follow-elem'
import ConfigRetrievalContent from '@/app/components/app/configuration/dataset-config/params-config/config-content'
import { RETRIEVE_TYPE } from '@/types/app'
import type {
DatasetConfigs,
} from '@/models/debug'
import { ChevronDown } from '@/app/components/base/icons/src/vender/line/arrows'
type Props = {
payload: {
retrieval_mode: RETRIEVE_TYPE
multiple_retrieval_config?: MultipleRetrievalConfig
}
onRetrievalModeChange: (mode: RETRIEVE_TYPE) => void
onMultipleRetrievalConfigChange: (config: MultipleRetrievalConfig) => void
}
const RetrievalConfig: FC<Props> = ({
payload,
onRetrievalModeChange,
onMultipleRetrievalConfigChange,
}) => {
const { t } = useTranslation()
const [open, setOpen] = useState(false)
const { multiple_retrieval_config } = payload
const handleChange = useCallback((configs: DatasetConfigs, isRetrievalModeChange?: boolean) => {
if (isRetrievalModeChange) {
onRetrievalModeChange(configs.retrieval_model)
return
}
onMultipleRetrievalConfigChange({
top_k: configs.top_k,
score_threshold: configs.score_threshold_enabled ? configs.score_threshold : null,
reranking_model: {
provider: configs.reranking_model?.reranking_provider_name,
model: configs.reranking_model?.reranking_model_name,
},
})
}, [onRetrievalModeChange, onMultipleRetrievalConfigChange])
return (
<PortalToFollowElem
open={open}
onOpenChange={setOpen}
placement='bottom-end'
offset={{
// mainAxis: 12,
crossAxis: -2,
}}
>
<PortalToFollowElemTrigger
onClick={() => setOpen(v => !v)}
>
<div className={cn(open && 'bg-gray-100', 'flex items-center h-6 px-2 rounded-md hover:bg-gray-100 group cursor-pointer select-none')}>
<div className={cn(open ? 'text-gray-700' : 'text-gray-500', 'leading-[18px] text-xs font-medium group-hover:bg-gray-100')}>{payload.retrieval_mode === RETRIEVE_TYPE.oneWay ? t('appDebug.datasetConfig.retrieveOneWay.title') : t('appDebug.datasetConfig.retrieveMultiWay.title')}</div>
<ChevronDown className='ml-1 w-3 h-3' />
</div>
</PortalToFollowElemTrigger>
<PortalToFollowElemContent style={{ zIndex: 1001 }}>
<div className='w-[404px] pt-3 pb-4 px-4 shadow-xl rounded-2xl border border-gray-200 bg-white'>
<ConfigRetrievalContent
datasetConfigs={
{
retrieval_model: payload.retrieval_mode,
reranking_model: {
reranking_provider_name: multiple_retrieval_config?.reranking_model?.provider || '',
reranking_model_name: multiple_retrieval_config?.reranking_model?.model || '',
},
top_k: multiple_retrieval_config?.top_k || 5,
score_threshold_enabled: !(multiple_retrieval_config?.score_threshold === undefined || multiple_retrieval_config?.score_threshold === null),
score_threshold: multiple_retrieval_config?.score_threshold || 0.5,
datasets: {
datasets: [],
},
}
}
onChange={handleChange}
/>
</div>
</PortalToFollowElemContent>
</PortalToFollowElem>
)
}
export default React.memo(RetrievalConfig)
...@@ -3,10 +3,10 @@ import { useTranslation } from 'react-i18next' ...@@ -3,10 +3,10 @@ import { useTranslation } from 'react-i18next'
import VarReferencePicker from '../_base/components/variable/var-reference-picker' import VarReferencePicker from '../_base/components/variable/var-reference-picker'
import useConfig from './use-config' import useConfig from './use-config'
import { mockData } from './mock' import { mockData } from './mock'
import RetrievalConfig from './components/retrieval-config'
import Field from '@/app/components/workflow/nodes/_base/components/field' import Field from '@/app/components/workflow/nodes/_base/components/field'
import Split from '@/app/components/workflow/nodes/_base/components/split' import Split from '@/app/components/workflow/nodes/_base/components/split'
import OutputVars, { VarItem } from '@/app/components/workflow/nodes/_base/components/output-vars' import OutputVars, { VarItem } from '@/app/components/workflow/nodes/_base/components/output-vars'
const i18nPrefix = 'workflow.nodes.knowledgeRetrieval' const i18nPrefix = 'workflow.nodes.knowledgeRetrieval'
const Panel: FC = () => { const Panel: FC = () => {
...@@ -16,6 +16,8 @@ const Panel: FC = () => { ...@@ -16,6 +16,8 @@ const Panel: FC = () => {
const { const {
inputs, inputs,
handleQueryVarChange, handleQueryVarChange,
handleRetrievalModeChange,
handleMultipleRetrievalConfigChange,
} = useConfig(mockData) } = useConfig(mockData)
return ( return (
...@@ -31,6 +33,24 @@ const Panel: FC = () => { ...@@ -31,6 +33,24 @@ const Panel: FC = () => {
onChange={handleQueryVarChange} onChange={handleQueryVarChange}
/> />
</Field> </Field>
<Field
title={t(`${i18nPrefix}.knowledge`)}
operations={
<div className='flex'>
<RetrievalConfig
payload={{
retrieval_mode: inputs.retrieval_mode,
multiple_retrieval_config: inputs.multiple_retrieval_config,
}}
onRetrievalModeChange={handleRetrievalModeChange}
onMultipleRetrievalConfigChange={handleMultipleRetrievalConfigChange}
/>
</div>
}
>
</Field>
</div> </div>
<Split /> <Split />
......
import type { CommonNodeType, ValueSelector } from '@/app/components/workflow/types' import type { CommonNodeType, ValueSelector } from '@/app/components/workflow/types'
import type { RETRIEVE_TYPE } from '@/types/app' import type { RETRIEVE_TYPE } from '@/types/app'
export type MultipleRetrievalConfig = {
top_k: number
score_threshold: number | null | undefined
reranking_model: {
provider: string
model: string
}
}
export type KnowledgeRetrievalNodeType = CommonNodeType & { export type KnowledgeRetrievalNodeType = CommonNodeType & {
query_variable_selector: ValueSelector query_variable_selector: ValueSelector
dataset_ids: string[] dataset_ids: string[]
retrieval_mode: RETRIEVE_TYPE retrieval_mode: RETRIEVE_TYPE
multiple_retrieval_config?: { multiple_retrieval_config?: MultipleRetrievalConfig
top_k: number
score_threshold: number
reranking_model: {
provider: string
model: string
}
}
} }
import { useCallback, useState } from 'react' import { useCallback, useState } from 'react'
import produce from 'immer' import produce from 'immer'
import type { ValueSelector } from '../../types' import type { ValueSelector } from '../../types'
import type { KnowledgeRetrievalNodeType } from './types' import type { KnowledgeRetrievalNodeType, MultipleRetrievalConfig } from './types'
import type { RETRIEVE_TYPE } from '@/types/app'
const useConfig = (initInputs: KnowledgeRetrievalNodeType) => { const useConfig = (initInputs: KnowledgeRetrievalNodeType) => {
const [inputs, setInputs] = useState<KnowledgeRetrievalNodeType>(initInputs) const [inputs, setInputs] = useState<KnowledgeRetrievalNodeType>(initInputs)
...@@ -13,9 +14,25 @@ const useConfig = (initInputs: KnowledgeRetrievalNodeType) => { ...@@ -13,9 +14,25 @@ const useConfig = (initInputs: KnowledgeRetrievalNodeType) => {
setInputs(newInputs) setInputs(newInputs)
}, [inputs, setInputs]) }, [inputs, setInputs])
const handleRetrievalModeChange = useCallback((newMode: RETRIEVE_TYPE) => {
const newInputs = produce(inputs, (draft) => {
draft.retrieval_mode = newMode
})
setInputs(newInputs)
}, [inputs, setInputs])
const handleMultipleRetrievalConfigChange = useCallback((newConfig: MultipleRetrievalConfig) => {
const newInputs = produce(inputs, (draft) => {
draft.multiple_retrieval_config = newConfig
})
setInputs(newInputs)
}, [inputs, setInputs])
return { return {
inputs, inputs,
handleQueryVarChange, handleQueryVarChange,
handleRetrievalModeChange,
handleMultipleRetrievalConfigChange,
} }
} }
......
...@@ -56,6 +56,7 @@ const translation = { ...@@ -56,6 +56,7 @@ const translation = {
}, },
knowledgeRetrieval: { knowledgeRetrieval: {
queryVariable: 'Query Variable', queryVariable: 'Query Variable',
knowledge: 'Knowledge',
outputVars: { outputVars: {
output: 'Retrieval segmented data', output: 'Retrieval segmented data',
content: 'Segmented content', content: 'Segmented content',
......
...@@ -55,6 +55,7 @@ const translation = { ...@@ -55,6 +55,7 @@ const translation = {
}, },
knowledgeRetrieval: { knowledgeRetrieval: {
queryVariable: '查询变量', queryVariable: '查询变量',
knowledge: '知识库',
outputVars: { outputVars: {
output: '召回的分段', output: '召回的分段',
content: '分段内容', content: '分段内容',
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment