import { useState, useEffect, useCallback } from 'react'; import { message } from 'antd'; import type { Dataset } from '~/api/dify-dataset/type/datasetTypes'; import type { Document } from '~/api/dify-dataset/type/documentTypes'; import { fetchDatasets } from '~/api/dify-dataset/api/datasetApi'; import { fetchDocuments } from '~/api/dify-dataset/api/documentApi'; import type { MenuTab } from '~/types/dify-dataset-manager/layout'; import { DEFAULT_DOCUMENT_PAGE_SIZE } from '~/types/dify-dataset-manager/index'; /** * 知识库管理器状态管理 Hook */ export function useDatasetManager() { // 知识库状态 const [dataset, setDataset] = useState(null); const [loadingDataset, setLoadingDataset] = useState(true); // 文档状态 const [documents, setDocuments] = useState([]); const [loadingDocuments, setLoadingDocuments] = useState(false); const [documentTotal, setDocumentTotal] = useState(0); const [documentPage, setDocumentPage] = useState(1); const [documentPageSize] = useState(DEFAULT_DOCUMENT_PAGE_SIZE); // 初始化状态 const [inited, setInited] = useState(false); const [error, setError] = useState(null); // 菜单状态 const [activeTab, setActiveTab] = useState('documents'); // 选中的文档(用于查看文档详情) const [selectedDocument, setSelectedDocument] = useState(null); /** * 加载文档列表 */ const loadDocuments = useCallback(async (datasetId: string, page: number = 1) => { if (!datasetId) return; setLoadingDocuments(true); try { console.log('[DatasetManager] 加载文档列表:', { datasetId, page }); const response = await fetchDocuments(datasetId, page, documentPageSize); console.log('[DatasetManager] 文档列表响应:', response); if (response && response.data) { setDocuments(response.data); setDocumentTotal(response.total); setDocumentPage(page); } } catch (err: any) { console.error('[DatasetManager] 加载文档列表失败:', err); message.error('加载文档列表失败'); } finally { setLoadingDocuments(false); } }, [documentPageSize]); /** * 加载知识库(获取第一个知识库) */ const loadDataset = useCallback(async () => { setLoadingDataset(true); try { console.log('[DatasetManager] 加载知识库...'); const response = await fetchDatasets(1, 1); console.log('[DatasetManager] 知识库响应:', response); if (response && response.data && response.data.length > 0) { const firstDataset = response.data[0]; setDataset(firstDataset); // 立即加载文档 await loadDocuments(firstDataset.id, 1); } else { setError('未找到知识库,请先在Dify中创建知识库'); } } catch (err: any) { console.error('[DatasetManager] 加载知识库失败:', err); setError(err.message || '加载知识库失败'); message.error('加载知识库失败'); } finally { setLoadingDataset(false); setInited(true); } }, [loadDocuments]); /** * 处理文档页码变化 */ const handlePageChange = useCallback((page: number) => { if (dataset) { loadDocuments(dataset.id, page); } }, [dataset, loadDocuments]); /** * 处理文档删除 */ const handleDocumentDeleted = useCallback((documentId: string) => { setDocuments((prev) => prev.filter((doc) => doc.id !== documentId)); setDocumentTotal((prev) => prev - 1); // 更新知识库的文档数量 setDataset((prev) => { if (prev) { return { ...prev, document_count: prev.document_count - 1 }; } return prev; }); }, []); /** * 处理文档状态变化 */ const handleDocumentStatusChanged = useCallback((documentId: string, enabled: boolean) => { setDocuments((prev) => prev.map((doc) => doc.id === documentId ? { ...doc, enabled } : doc ) ); }, []); /** * 刷新文档列表 */ const handleRefresh = useCallback(() => { if (dataset) { loadDocuments(dataset.id, documentPage); } }, [dataset, documentPage, loadDocuments]); /** * 查看文档详情(分段管理) */ const handleViewDocument = useCallback((doc: Document) => { console.log('[DatasetManager] 查看文档详情:', doc); setSelectedDocument(doc); }, []); /** * 返回文档列表 */ const handleBackToDocuments = useCallback(() => { setSelectedDocument(null); }, []); /** * 处理菜单切换 */ const handleTabChange = useCallback((tab: MenuTab) => { setActiveTab(tab); // 切换菜单时清除选中的文档 if (tab !== 'documents') { setSelectedDocument(null); } }, []); /** * 处理知识库更新 */ const handleDatasetUpdated = useCallback((updatedDataset: Dataset) => { setDataset(updatedDataset); }, []); // 初始化 useEffect(() => { loadDataset(); }, [loadDataset]); return { // 状态 dataset, loadingDataset, documents, loadingDocuments, documentTotal, documentPage, documentPageSize, inited, error, activeTab, selectedDocument, // 方法 loadDataset, loadDocuments, handlePageChange, handleDocumentDeleted, handleDocumentStatusChanged, handleRefresh, handleViewDocument, handleBackToDocuments, handleTabChange, handleDatasetUpdated, }; } export type UseDatasetManagerReturn = ReturnType;