'use client' import { useCallback, useRef, useState } from 'react' import { useMount } from 'ahooks' import { useTranslation } from 'react-i18next' import PermissionSelector from '../permission-selector' import IndexMethod from '../index-method' import RetrievalSettings from '../../external-knowledge-base/create/RetrievalSettings' import { IndexingType } from '../../create/step-two' import RetrievalMethodConfig from '@/app/components/datasets/common/retrieval-method-config' import EconomicalRetrievalMethodConfig from '@/app/components/datasets/common/economical-retrieval-method-config' import Button from '@/app/components/base/button' import Input from '@/app/components/base/input' import Textarea from '@/app/components/base/textarea' import { ApiConnectionMod } from '@/app/components/base/icons/src/vender/solid/development' import { updateDatasetSetting } from '@/service/datasets' import type { IconInfo } from '@/models/datasets' import { ChunkingMode, DatasetPermission } from '@/models/datasets' import { useDatasetDetailContextWithSelector } from '@/context/dataset-detail' import type { AppIconType, RetrievalConfig } from '@/types/app' import { useSelector as useAppContextWithSelector } from '@/context/app-context' import { isReRankModelSelected } from '@/app/components/datasets/common/check-rerank-model' import ModelSelector from '@/app/components/header/account-setting/model-provider-page/model-selector' import { useModelList, useModelListAndDefaultModelAndCurrentProviderAndModel, } from '@/app/components/header/account-setting/model-provider-page/hooks' import type { DefaultModel } from '@/app/components/header/account-setting/model-provider-page/declarations' import { ModelTypeEnum } from '@/app/components/header/account-setting/model-provider-page/declarations' import { fetchMembers } from '@/service/common' import type { Member } from '@/models/common' import AppIcon from '@/app/components/base/app-icon' import type { AppIconSelection } from '@/app/components/base/app-icon-picker' import AppIconPicker from '@/app/components/base/app-icon-picker' import Divider from '@/app/components/base/divider' import ChunkStructure from '../chunk-structure' import Toast from '@/app/components/base/toast' import { RiAlertFill } from '@remixicon/react' import { useDocLink } from '@/context/i18n' import { useResetDatasetList } from '@/service/knowledge/use-dataset' const rowClass = 'flex gap-x-1' const labelClass = 'flex items-center shrink-0 w-[180px] h-7 pt-1' const DEFAULT_APP_ICON: IconInfo = { icon_type: 'emoji', icon: '📙', icon_background: '#FFF4ED', icon_url: '', } const Form = () => { const { t } = useTranslation() const docLink = useDocLink() const isCurrentWorkspaceDatasetOperator = useAppContextWithSelector(state => state.isCurrentWorkspaceDatasetOperator) const currentDataset = useDatasetDetailContextWithSelector(state => state.dataset) const mutateDatasets = useDatasetDetailContextWithSelector(state => state.mutateDatasetRes) const [loading, setLoading] = useState(false) const [name, setName] = useState(currentDataset?.name ?? '') const [iconInfo, setIconInfo] = useState(currentDataset?.icon_info || DEFAULT_APP_ICON) const [showAppIconPicker, setShowAppIconPicker] = useState(false) const [description, setDescription] = useState(currentDataset?.description ?? '') const [permission, setPermission] = useState(currentDataset?.permission) const [topK, setTopK] = useState(currentDataset?.external_retrieval_model.top_k ?? 2) const [scoreThreshold, setScoreThreshold] = useState(currentDataset?.external_retrieval_model.score_threshold ?? 0.5) const [scoreThresholdEnabled, setScoreThresholdEnabled] = useState(currentDataset?.external_retrieval_model.score_threshold_enabled ?? false) const [selectedMemberIDs, setSelectedMemberIDs] = useState(currentDataset?.partial_member_list || []) const [memberList, setMemberList] = useState([]) const [indexMethod, setIndexMethod] = useState(currentDataset?.indexing_technique) const [keywordNumber, setKeywordNumber] = useState(currentDataset?.keyword_number ?? 10) const [retrievalConfig, setRetrievalConfig] = useState(currentDataset?.retrieval_model_dict as RetrievalConfig) const [embeddingModel, setEmbeddingModel] = useState( currentDataset?.embedding_model ? { provider: currentDataset.embedding_model_provider, model: currentDataset.embedding_model, } : { provider: '', model: '', }, ) const { modelList: rerankModelList, } = useModelListAndDefaultModelAndCurrentProviderAndModel(ModelTypeEnum.rerank) const { data: embeddingModelList } = useModelList(ModelTypeEnum.textEmbedding) const previousAppIcon = useRef(DEFAULT_APP_ICON) const getMembers = async () => { const { accounts } = await fetchMembers({ url: '/workspaces/current/members', params: {} }) if (!accounts) setMemberList([]) else setMemberList(accounts) } const handleOpenAppIconPicker = useCallback(() => { setShowAppIconPicker(true) previousAppIcon.current = iconInfo }, [iconInfo]) const handleSelectAppIcon = useCallback((icon: AppIconSelection) => { const iconInfo: IconInfo = { icon_type: icon.type, icon: icon.type === 'emoji' ? icon.icon : icon.fileId, icon_background: icon.type === 'emoji' ? icon.background : undefined, icon_url: icon.type === 'emoji' ? undefined : icon.url, } setIconInfo(iconInfo) setShowAppIconPicker(false) }, []) const handleCloseAppIconPicker = useCallback(() => { setIconInfo(previousAppIcon.current) setShowAppIconPicker(false) }, []) const handleSettingsChange = useCallback((data: { top_k?: number; score_threshold?: number; score_threshold_enabled?: boolean }) => { if (data.top_k !== undefined) setTopK(data.top_k) if (data.score_threshold !== undefined) setScoreThreshold(data.score_threshold) if (data.score_threshold_enabled !== undefined) setScoreThresholdEnabled(data.score_threshold_enabled) }, []) useMount(() => { getMembers() }) const resetDatasetList = useResetDatasetList() const handleSave = async () => { if (loading) return if (!name?.trim()) { Toast.notify({ type: 'error', message: t('datasetSettings.form.nameError') }) return } if ( !isReRankModelSelected({ rerankModelList, retrievalConfig, indexMethod, }) ) { Toast.notify({ type: 'error', message: t('appDebug.datasetConfig.rerankModelRequired') }) return } if (retrievalConfig.weights) { retrievalConfig.weights.vector_setting.embedding_provider_name = currentDataset?.embedding_model_provider || '' retrievalConfig.weights.vector_setting.embedding_model_name = currentDataset?.embedding_model || '' } try { setLoading(true) const requestParams = { datasetId: currentDataset!.id, body: { name, icon_info: iconInfo, doc_form: currentDataset?.doc_form, description, permission, indexing_technique: indexMethod, retrieval_model: { ...retrievalConfig, score_threshold: retrievalConfig.score_threshold_enabled ? retrievalConfig.score_threshold : 0, }, embedding_model: embeddingModel.model, embedding_model_provider: embeddingModel.provider, ...(currentDataset!.provider === 'external' && { external_knowledge_id: currentDataset!.external_knowledge_info.external_knowledge_id, external_knowledge_api_id: currentDataset!.external_knowledge_info.external_knowledge_api_id, external_retrieval_model: { top_k: topK, score_threshold: scoreThreshold, score_threshold_enabled: scoreThresholdEnabled, }, }), keyword_number: keywordNumber, }, } as any if (permission === DatasetPermission.partialMembers) { requestParams.body.partial_member_list = selectedMemberIDs.map((id) => { return { user_id: id, role: memberList.find(member => member.id === id)?.role, } }) } await updateDatasetSetting(requestParams) Toast.notify({ type: 'success', message: t('common.actionMsg.modifiedSuccessfully') }) if (mutateDatasets) { await mutateDatasets() resetDatasetList() } } catch { Toast.notify({ type: 'error', message: t('common.actionMsg.modifiedUnsuccessfully') }) } finally { setLoading(false) } } const isShowIndexMethod = currentDataset && currentDataset.doc_form !== ChunkingMode.parentChild && currentDataset.indexing_technique && indexMethod return (
{/* Dataset name and icon */}
{t('datasetSettings.form.nameAndIcon')}
setName(e.target.value)} />
{/* Dataset description */}
{t('datasetSettings.form.desc')}