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, fetchDataset } from '~/api/dify-dataset/api/datasetApi'; import { fetchDocuments } from '~/api/dify-dataset/api/documentApi'; import { getMyDatasets, type AreaDataset } from '~/api/v3/dify/area-datasets'; 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 [availableDatasets, setAvailableDatasets] = useState([]); const [loadingAvailableDatasets, setLoadingAvailableDatasets] = 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 loadAvailableDatasets = useCallback(async () => { setLoadingAvailableDatasets(true); try { console.log('[DatasetManager] 加载用户可访问的知识库列表...'); const response = await getMyDatasets(); console.log('[DatasetManager] 用户知识库列表响应:', response); if (response && response.code === 0 && response.data) { const dataList = Array.isArray(response.data.data) ? response.data.data : []; setAvailableDatasets(dataList); return dataList; } else { console.error('[DatasetManager] 获取用户知识库列表失败:', response); setAvailableDatasets([]); return []; } } catch (err: any) { console.error('[DatasetManager] 加载用户知识库列表失败:', err); setAvailableDatasets([]); return []; } finally { setLoadingAvailableDatasets(false); } }, []); /** * 根据 dataset_id 加载知识库详情 */ const loadDatasetById = useCallback(async (datasetId: string) => { setLoadingDataset(true); try { console.log('[DatasetManager] 加载知识库详情:', datasetId); const fullDataset = await fetchDataset(datasetId); console.log('[DatasetManager] 知识库详情响应:', fullDataset); setDataset(fullDataset); // 立即加载文档 await loadDocuments(datasetId, 1); } catch (err: any) { console.error('[DatasetManager] 加载知识库详情失败:', err); const is403 = err.message?.includes('403') || err.response?.status === 403; const msg = is403 ? '您没有查看知识库的权限' : (err.message || '加载知识库失败'); setError(msg); message.error(msg); } finally { setLoadingDataset(false); } }, [loadDocuments]); /** * 加载知识库(获取第一个知识库,再获取详情以包含 retrieval_model) */ const loadDataset = useCallback(async () => { setLoadingDataset(true); try { console.log('[DatasetManager] 加载知识库...'); // 先加载用户可访问的知识库列表 const userDatasets = await loadAvailableDatasets(); if (userDatasets.length > 0) { // 找到默认知识库或第一个知识库 const defaultDataset = userDatasets.find(ds => ds.is_default) || userDatasets[0]; const datasetId = defaultDataset.dataset_id; console.log('[DatasetManager] 使用知识库:', defaultDataset.dataset_name, datasetId); // 获取知识库详情 const fullDataset = await fetchDataset(datasetId); console.log('[DatasetManager] 知识库详情响应:', fullDataset); setDataset(fullDataset); // 立即加载文档 await loadDocuments(datasetId, 1); } else { // 回退到原有逻辑:直接从 Dify 获取 console.log('[DatasetManager] 用户无绑定知识库,使用默认逻辑...'); const response = await fetchDatasets(1, 1); console.log('[DatasetManager] Dify知识库列表响应:', response); if (response && response.data && response.data.length > 0) { const firstDatasetId = response.data[0].id; const fullDataset = await fetchDataset(firstDatasetId); console.log('[DatasetManager] 知识库详情响应:', fullDataset); setDataset(fullDataset); await loadDocuments(firstDatasetId, 1); } else { setError('未找到知识库,请先联系管理员创建知识库'); } } } catch (err: any) { console.error('[DatasetManager] 加载知识库失败:', err); const is403 = err.message?.includes('403') || err.response?.status === 403; const msg = is403 ? '您没有查看知识库的权限' : (err.message || '加载知识库失败'); setError(msg); message.error(msg); } finally { setLoadingDataset(false); setInited(true); } }, [loadDocuments, loadAvailableDatasets]); /** * 切换知识库 */ const handleDatasetChange = useCallback(async (datasetId: string) => { console.log('[DatasetManager] 切换知识库:', datasetId); // 重置状态 setSelectedDocument(null); setActiveTab('documents'); // 加载新知识库 await loadDatasetById(datasetId); }, [loadDatasetById]); /** * 处理文档页码变化 */ 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, // 知识库列表(基于权限) availableDatasets, loadingAvailableDatasets, // 方法 loadDataset, loadDocuments, handlePageChange, handleDocumentDeleted, handleDocumentStatusChanged, handleRefresh, handleViewDocument, handleBackToDocuments, handleTabChange, handleDatasetUpdated, handleDatasetChange, }; } export type UseDatasetManagerReturn = ReturnType;