feat: retrieval config

2024-02-29 18:24:05 +08:00 · 2024-02-29 18:24:05 +08:00 · 257e795ca9
parent bafdc510d6
commit 257e795ca9
8 changed files with 292 additions and 117 deletions
--- a/web/app/components/app/configuration/dataset-config/params-config/config-content.tsx
+++ b/web/app/components/app/configuration/dataset-config/params-config/config-content.tsx
@ -0,0 +1,139 @@
+'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)
--- a/web/app/components/app/configuration/dataset-config/params-config/index.tsx
+++ b/web/app/components/app/configuration/dataset-config/params-config/index.tsx
@ -4,21 +4,14 @@ import { memo, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import { useContext } from 'use-context-selector'
 import cn from 'classnames'
+import ConfigContent from './config-content'
 import { Settings04 } from '@/app/components/base/icons/src/vender/line/general'
 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 Button from '@/app/components/base/button'
-import RadioCard from '@/app/components/base/radio-card/simple'
 import { RETRIEVE_TYPE } from '@/types/app'
 import Toast from '@/app/components/base/toast'
 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'

 const ParamsConfig: FC = () => {
@ -30,58 +23,11 @@ const ParamsConfig: FC = () => {
  } = useContext(ConfigContext)
  const [tempDataSetConfigs, setTempDataSetConfigs] = useState(datasetConfigs)

-  const type = tempDataSetConfigs.retrieval_model
-  const setType = (value: RETRIEVE_TYPE) => {
-    setTempDataSetConfigs({
-      ...tempDataSetConfigs,
-      retrieval_model: value,
-    })
-  }
  const {
-    modelList: rerankModelList,
    defaultModel: rerankDefaultModel,
    currentModel: isRerankDefaultModelVaild,
  } = 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 = () => {
    let errMsg = ''
    if (tempDataSetConfigs.retrieval_model === RETRIEVE_TYPE.multiWay) {
@ -140,58 +86,11 @@ const ParamsConfig: FC = () => {
            wrapperClassName='z-50'
            title={t('appDebug.datasetConfig.settingTitle')}
          >
-            <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) => {
-                        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>
-              </>
-            )}
+            <ConfigContent
+              datasetConfigs={tempDataSetConfigs}
+              onChange={setTempDataSetConfigs}
+            />
+
            <div className='mt-6 flex justify-end'>
              <Button className='mr-2 flex-shrink-0' onClick={() => {
                setOpen(false)
--- a/web/app/components/workflow/nodes/knowledge-retrieval/components/retrieval-config.tsx
+++ b/web/app/components/workflow/nodes/knowledge-retrieval/components/retrieval-config.tsx
@ -0,0 +1,97 @@
+'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)
--- a/web/app/components/workflow/nodes/knowledge-retrieval/panel.tsx
+++ b/web/app/components/workflow/nodes/knowledge-retrieval/panel.tsx
@ -3,10 +3,10 @@ import { useTranslation } from 'react-i18next'
 import VarReferencePicker from '../_base/components/variable/var-reference-picker'
 import useConfig from './use-config'
 import { mockData } from './mock'
+import RetrievalConfig from './components/retrieval-config'
 import Field from '@/app/components/workflow/nodes/_base/components/field'
 import Split from '@/app/components/workflow/nodes/_base/components/split'
 import OutputVars, { VarItem } from '@/app/components/workflow/nodes/_base/components/output-vars'
-
 const i18nPrefix = 'workflow.nodes.knowledgeRetrieval'

 const Panel: FC = () => {
@ -16,6 +16,8 @@ const Panel: FC = () => {
  const {
    inputs,
    handleQueryVarChange,
+    handleRetrievalModeChange,
+    handleMultipleRetrievalConfigChange,
  } = useConfig(mockData)

  return (
@ -31,6 +33,24 @@ const Panel: FC = () => {
            onChange={handleQueryVarChange}
          />
        </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>

      <Split />
--- a/web/app/components/workflow/nodes/knowledge-retrieval/types.ts
+++ b/web/app/components/workflow/nodes/knowledge-retrieval/types.ts
@ -1,16 +1,17 @@
 import type { CommonNodeType, ValueSelector } from '@/app/components/workflow/types'
 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 & {
  query_variable_selector: ValueSelector
  dataset_ids: string[]
  retrieval_mode: RETRIEVE_TYPE
-  multiple_retrieval_config?: {
-    top_k: number
-    score_threshold: number
-    reranking_model: {
-      provider: string
-      model: string
-    }
-  }
+  multiple_retrieval_config?: MultipleRetrievalConfig
 }
--- a/web/app/components/workflow/nodes/knowledge-retrieval/use-config.ts
+++ b/web/app/components/workflow/nodes/knowledge-retrieval/use-config.ts
@ -1,7 +1,8 @@
 import { useCallback, useState } from 'react'
 import produce from 'immer'
 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 [inputs, setInputs] = useState<KnowledgeRetrievalNodeType>(initInputs)
@ -13,9 +14,25 @@ const useConfig = (initInputs: KnowledgeRetrievalNodeType) => {
    setInputs(newInputs)
  }, [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 {
    inputs,
    handleQueryVarChange,
+    handleRetrievalModeChange,
+    handleMultipleRetrievalConfigChange,
  }
 }

--- a/web/i18n/en-US/workflow.ts
+++ b/web/i18n/en-US/workflow.ts
@ -56,6 +56,7 @@ const translation = {
    },
    knowledgeRetrieval: {
      queryVariable: 'Query Variable',
+      knowledge: 'Knowledge',
      outputVars: {
        output: 'Retrieval segmented data',
        content: 'Segmented content',
--- a/web/i18n/zh-Hans/workflow.ts
+++ b/web/i18n/zh-Hans/workflow.ts
@ -55,6 +55,7 @@ const translation = {
    },
    knowledgeRetrieval: {
      queryVariable: '查询变量',
+      knowledge: '知识库',
      outputVars: {
        output: '召回的分段',
        content: '分段内容',