import { memo, useRef, useState } from 'react' import type { FC } from 'react' import { useTranslation } from 'react-i18next' import { useContext } from 'use-context-selector' import { useParams } from 'next/navigation' import { RiCloseLine, RiExpandDiagonalLine } from '@remixicon/react' import { useShallow } from 'zustand/react/shallow' import { useSegmentListContext } from './completed' import { SegmentIndexTag } from './completed/common/segment-index-tag' import ActionButtons from './completed/common/action-buttons' import Keywords from './completed/common/keywords' import ChunkContent from './completed/common/chunk-content' import AddAnother from './completed/common/add-another' import { useDocumentContext } from './index' import { useStore as useAppStore } from '@/app/components/app/store' import { ToastContext } from '@/app/components/base/toast' import type { SegmentUpdater } from '@/models/datasets' import { addSegment } from '@/service/datasets' import classNames from '@/utils/classnames' import { formatNumber } from '@/utils/format' import Divider from '@/app/components/base/divider' type NewSegmentModalProps = { onCancel: () => void docForm: string onSave: () => void viewNewlyAddedChunk: () => void } const NewSegmentModal: FC = ({ onCancel, docForm, onSave, viewNewlyAddedChunk, }) => { const { t } = useTranslation() const { notify } = useContext(ToastContext) const [question, setQuestion] = useState('') const [answer, setAnswer] = useState('') const { datasetId, documentId } = useParams<{ datasetId: string; documentId: string }>() const [keywords, setKeywords] = useState([]) const [loading, setLoading] = useState(false) const [addAnother, setAddAnother] = useState(true) const [fullScreen, toggleFullScreen] = useSegmentListContext(s => [s.fullScreen, s.toggleFullScreen]) const mode = useDocumentContext(s => s.mode) const { appSidebarExpand } = useAppStore(useShallow(state => ({ appSidebarExpand: state.appSidebarExpand, }))) const refreshTimer = useRef(null) const CustomButton = <> const handleCancel = (actionType: 'esc' | 'add' = 'esc') => { if (actionType === 'esc' || !addAnother) onCancel() setQuestion('') setAnswer('') setKeywords([]) } const handleSave = async () => { const params: SegmentUpdater = { content: '' } if (docForm === 'qa_model') { if (!question.trim()) return notify({ type: 'error', message: t('datasetDocuments.segment.questionEmpty') }) if (!answer.trim()) return notify({ type: 'error', message: t('datasetDocuments.segment.answerEmpty') }) params.content = question params.answer = answer } else { if (!question.trim()) return notify({ type: 'error', message: t('datasetDocuments.segment.contentEmpty') }) params.content = question } if (keywords?.length) params.keywords = keywords setLoading(true) try { await addSegment({ datasetId, documentId, body: params }) notify({ type: 'success', message: t('datasetDocuments.segment.chunkAdded'), className: `!w-[296px] !bottom-0 ${appSidebarExpand === 'expand' ? '!left-[216px]' : '!left-14'} !top-auto !right-auto !mb-[52px] !ml-11`, customComponent: CustomButton, }) handleCancel('add') refreshTimer.current = setTimeout(() => { onSave() }, 3000) } finally { setLoading(false) } } return (
{ docForm === 'qa_model' ? t('datasetDocuments.segment.newQaSegment') : t('datasetDocuments.segment.addChunk') }
ยท {formatNumber(question.length)} {t('datasetDocuments.segment.characters')}
{fullScreen && ( <> setAddAnother(!addAnother)} /> )}
setQuestion(question)} onAnswerChange={answer => setAnswer(answer)} isEditMode={true} />
{mode === 'custom' && setKeywords(keywords)} />}
{!fullScreen && (
setAddAnother(!addAnother)} />
)}
) } export default memo(NewSegmentModal)