diff --git a/web/src/components/knowledge-base-item.tsx b/web/src/components/knowledge-base-item.tsx index a2dac7cbb1..ed970179b6 100644 --- a/web/src/components/knowledge-base-item.tsx +++ b/web/src/components/knowledge-base-item.tsx @@ -1,5 +1,8 @@ import { DocumentParserType } from '@/constants/knowledge'; -import { useFetchKnowledgeList } from '@/hooks/use-knowledge-request'; +import { + useFetchDatasetsByIds, + useFetchKnowledgeList, +} from '@/hooks/use-knowledge-request'; import { IDataset } from '@/interfaces/database/dataset'; import { useBuildQueryVariableOptions } from '@/pages/agent/hooks/use-get-begin-query'; import { useDebounce } from 'ahooks'; @@ -36,12 +39,31 @@ export function useDisableDifferenceEmbeddingDataset(name: string) { } = useFetchKnowledgeList(false, debouncedSearchString); const datasetCacheRef = useRef(new Map()); + const selectedDatasetIds = useMemo( + () => (Array.isArray(datasetId) ? datasetId : []), + [datasetId], + ); + + // Selected dataset IDs that are neither in the currently loaded page nor + // already cached. These need to be fetched by ID so their names can be + // echoed back in the form field (the paginated list may not contain them). + const missingIds = useMemo(() => { + const loadedIds = new Set(datasetListOrigin.map((d) => d.id)); + return selectedDatasetIds.filter( + (id) => !loadedIds.has(id) && !datasetCacheRef.current.has(id), + ); + }, [datasetListOrigin, selectedDatasetIds]); + + const { data: missingDatasets } = useFetchDatasetsByIds(missingIds); + const datasetList = useMemo(() => { datasetListOrigin.forEach((dataset) => { datasetCacheRef.current.set(dataset.id, dataset); }); + missingDatasets?.forEach((dataset) => { + datasetCacheRef.current.set(dataset.id, dataset); + }); - const selectedDatasetIds = Array.isArray(datasetId) ? datasetId : []; const selectedDatasets = selectedDatasetIds .map((id) => datasetCacheRef.current.get(id)) .filter(Boolean) as IDataset[]; @@ -54,7 +76,7 @@ export function useDisableDifferenceEmbeddingDataset(name: string) { ]), ).values(), ); - }, [datasetId, datasetListOrigin]); + }, [datasetListOrigin, selectedDatasetIds, missingDatasets]); const selectedEmbedId = useMemo(() => { const data = datasetList?.find((item) => item.id === datasetId?.[0]); diff --git a/web/src/hooks/use-knowledge-request.ts b/web/src/hooks/use-knowledge-request.ts index 4c336d2206..138f05f19f 100644 --- a/web/src/hooks/use-knowledge-request.ts +++ b/web/src/hooks/use-knowledge-request.ts @@ -41,6 +41,7 @@ import kbService, { listArtifacts, datasetFilter, listDataset, + listDatasetByIds, listTag, listWikiCommits, removeTag, @@ -1003,6 +1004,8 @@ export const KnowledgeListKeys = { keywords, pageSize, ] as const, + byIds: (ids: string[]) => + [KnowledgeApiAction.FetchKnowledgeList, 'byIds', ids] as const, }; export const useFetchKnowledgeList = ( @@ -1114,6 +1117,26 @@ export const useSelectKnowledgeOptions = () => { return options; }; +/** + * Fetch datasets by a set of IDs. Used to resolve the names of + * already-selected datasets that are not present in the first page of + * the paginated list so they can be echoed back in the form field. + */ +export const useFetchDatasetsByIds = (ids: string[]) => { + const sortedIds = useMemo(() => [...ids].sort(), [ids]); + const { data, isFetching: loading } = useQuery({ + queryKey: KnowledgeListKeys.byIds(sortedIds), + enabled: sortedIds.length > 0, + gcTime: 0, + queryFn: async () => { + const { data } = await listDatasetByIds(sortedIds); + return (data?.data ?? []) as IDataset[]; + }, + }); + + return { data, loading }; +}; + //#region tags export const useRenameTag = () => { const knowledgeBaseId = useKnowledgeBaseId(); diff --git a/web/src/services/knowledge-service.ts b/web/src/services/knowledge-service.ts index c23f64c040..62dabd62ba 100644 --- a/web/src/services/knowledge-service.ts +++ b/web/src/services/knowledge-service.ts @@ -268,6 +268,14 @@ export function deleteKnowledgeGraph(knowledgeId: string) { export const listDataset = (params?: IFetchKnowledgeListRequestParams) => request.get(api.kbList, { params }); +// Fetch datasets by a set of IDs via the `ids` query param (comma-joined). +// Used to echo back already-selected datasets whose names are not present +// in the first page of the paginated list. +export const listDatasetByIds = (ids: string[]) => + request.get(api.kbList, { + params: { ids: ids.join(','), page_size: ids.length }, + }); + export const datasetFilter = () => request.get(api.datasetFilter); export const updateKb = (datasetId: string, data: Record) =>