import { memo, useMemo, 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 Dot from './completed/common/dot' import { useDocumentContext } from './index' import { useStore as useAppStore } from '@/app/components/app/store' import { ToastContext } from '@/app/components/base/toast' import { ChunkingMode, type SegmentUpdater } from '@/models/datasets' import classNames from '@/utils/classnames' import { formatNumber } from '@/utils/format' import Divider from '@/app/components/base/divider' import { useAddSegment } from '@/service/knowledge/use-segment' type NewSegmentModalProps = { onCancel: () => void docForm: ChunkingMode 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 = useSegmentListContext(s => s.fullScreen) const toggleFullScreen = useSegmentListContext(s => s.toggleFullScreen) const mode = useDocumentContext(s => s.mode) const { appSidebarExpand } = useAppStore(useShallow(state => ({ appSidebarExpand: state.appSidebarExpand, }))) const refreshTimer = useRef(null) const CustomButton = <> const isQAModel = useMemo(() => { return docForm === ChunkingMode.qa }, [docForm]) const handleCancel = (actionType: 'esc' | 'add' = 'esc') => { if (actionType === 'esc' || !addAnother) onCancel() } const { mutateAsync: addSegment } = useAddSegment() const handleSave = async () => { const params: SegmentUpdater = { content: '' } if (isQAModel) { 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) await addSegment({ datasetId, documentId, body: params }, { onSuccess() { 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) }, onSettled() { setLoading(false) }, }) } const wordCountText = useMemo(() => { const count = isQAModel ? (question.length + answer.length) : question.length return `${formatNumber(count)} ${t('datasetDocuments.segment.characters', { count })}` // eslint-disable-next-line react-hooks/exhaustive-deps }, [question.length, answer.length, isQAModel]) return (
{ t('datasetDocuments.segment.addChunk') }
{wordCountText}
{fullScreen && ( <> setAddAnother(!addAnother)} /> )}
setQuestion(question)} onAnswerChange={answer => setAnswer(answer)} isEditMode={true} />
{mode === 'custom' && setKeywords(keywords)} />}
{!fullScreen && (
setAddAnother(!addAnother)} />
)}
) } export default memo(NewSegmentModal)