import { message } from 'antd'; import { useCallback, useEffect, useRef, useState } from 'react'; import { fetchUploadFileInfo, downloadOriginalFile, updateDocumentByFile, fetchIndexingStatus, } from '~/api/dify-dataset/api/documentApi'; import { fetchSegments } from '~/api/dify-dataset/api/segmentApi'; import type { Segment, IndexingStatus } from '~/api/dify-dataset/type'; import type { Document } from '~/api/dify-dataset/type/documentTypes'; import type { DocumentDetailSegmentationSettings } from '~/types/dify-dataset-manager/document-detail'; import { DEFAULT_DOCUMENT_DETAIL_SETTINGS } from '~/types/dify-dataset-manager/document-detail'; /** * 文档详情状态管理 Hook */ export function useDocumentDetail(datasetId: string, document: Document | null) { // 分段设置状态 const [settings, setSettings] = useState(DEFAULT_DOCUMENT_DETAIL_SETTINGS); // 预览状态 const [previewSegments, setPreviewSegments] = useState([]); const [previewLoading, setPreviewLoading] = useState(false); const [showPreview, setShowPreview] = useState(false); // 保存状态 const [saving, setSaving] = useState(false); // 处理状态(嵌入进度) const [isProcessing, setIsProcessing] = useState(false); const [indexingStatus, setIndexingStatus] = useState(null); const pollingTimerRef = useRef | null>(null); /** * 停止轮询 */ const stopPolling = useCallback(() => { if (pollingTimerRef.current) { clearInterval(pollingTimerRef.current); pollingTimerRef.current = null; } }, []); /** * 轮询索引状态 */ const pollIndexingStatus = useCallback(async (batch: string) => { try { const response = await fetchIndexingStatus(datasetId, batch); const status = response.data?.[0]; if (status) { setIndexingStatus(status.indexing_status); if (status.indexing_status === 'completed') { // 停止轮询 stopPolling(); setIsProcessing(false); message.success('文档处理完成'); // 刷新分段预览 setPreviewLoading(true); try { const segmentResponse = await fetchSegments(datasetId, document?.id || '', 1, 50); setPreviewSegments(segmentResponse.data || []); setShowPreview(true); } catch (err) { console.error('刷新分段失败:', err); } finally { setPreviewLoading(false); } } else if (status.indexing_status === 'error') { // 停止轮询 stopPolling(); setIsProcessing(false); message.error(status.error || '处理失败'); } } } catch (err) { console.error('获取索引状态失败:', err); } }, [datasetId, document?.id, stopPolling]); /** * 开始轮询 */ const startPolling = useCallback((batch: string) => { stopPolling(); pollingTimerRef.current = setInterval(() => { pollIndexingStatus(batch); }, 2000); // 立即执行一次 pollIndexingStatus(batch); }, [stopPolling, pollIndexingStatus]); // 当文档变化时重置设置 useEffect(() => { if (document) { // 可以从文档中读取已有的设置,这里使用默认值 setSettings(DEFAULT_DOCUMENT_DETAIL_SETTINGS); setPreviewSegments([]); setShowPreview(false); setIsProcessing(false); setIndexingStatus(null); stopPolling(); } }, [document?.id, stopPolling]); // 组件卸载时清理定时器 useEffect(() => { return () => { stopPolling(); }; }, [stopPolling]); /** * 更新设置 */ const updateSettings = useCallback((key: keyof DocumentDetailSegmentationSettings, value: any) => { setSettings(prev => ({ ...prev, [key]: value })); }, []); /** * 重置设置 */ const handleReset = useCallback(() => { setSettings(DEFAULT_DOCUMENT_DETAIL_SETTINGS); setPreviewSegments([]); setShowPreview(false); }, []); /** * 预览分段 */ const handlePreview = useCallback(async () => { if (!document) return; setPreviewLoading(true); setShowPreview(true); try { // 获取当前文档的分段作为预览 const response = await fetchSegments(datasetId, document.id, 1, 50); setPreviewSegments(response.data || []); if (response.data?.length === 0) { message.info('该文档暂无分段数据'); } } catch (err: any) { console.error('预览分段失败:', err); message.error(err.message || '预览失败'); } finally { setPreviewLoading(false); } }, [datasetId, document]); /** * 保存并处理 * 流程:获取原始文件 → 下载 → 用新参数重新上传 → 轮询嵌入状态 */ const handleSaveAndProcess = useCallback(async () => { if (!document) return; setIsProcessing(true); setSaving(true); setIndexingStatus('waiting'); try { // 1. 获取原始文件信息 message.loading({ content: '正在获取原始文件信息...', key: 'save-process' }); const uploadFileInfo = await fetchUploadFileInfo(datasetId, document.id); // 2. 下载原始文件(通过代理路由) message.loading({ content: '正在下载原始文件...', key: 'save-process' }); const file = await downloadOriginalFile(uploadFileInfo); // 3. 用新参数重新上传 message.loading({ content: '正在应用新配置并重新处理...', key: 'save-process' }); const result = await updateDocumentByFile(datasetId, document.id, file, { indexing_technique: settings.indexingTechnique, process_rule: { mode: 'custom', rules: { pre_processing_rules: [ { id: 'remove_extra_spaces', enabled: settings.removeExtraSpaces }, { id: 'remove_urls_emails', enabled: settings.removeUrlsEmails }, ], segmentation: { separator: settings.separator.replace(/\\n/g, '\n'), max_tokens: settings.maxTokens, chunk_overlap: settings.chunkOverlap, }, }, }, }); message.success({ content: '文档正在处理中...', key: 'save-process' }); // 4. 开始轮询嵌入状态 startPolling(result.batch); } catch (err: any) { console.error('保存设置失败:', err); message.error({ content: err.message || '保存失败', key: 'save-process' }); setIsProcessing(false); setIndexingStatus(null); } finally { setSaving(false); } }, [datasetId, document, settings, startPolling]); return { // 状态 settings, previewSegments, previewLoading, showPreview, saving, isProcessing, indexingStatus, // 方法 updateSettings, handleReset, handlePreview, handleSaveAndProcess, }; } export type UseDocumentDetailReturn = ReturnType;