diff --git a/packages/toolkit/src/lib/react-query-service/knowledge/index.ts b/packages/toolkit/src/lib/react-query-service/knowledge/index.ts index d9d7f91753..f3d3c7a70c 100644 --- a/packages/toolkit/src/lib/react-query-service/knowledge/index.ts +++ b/packages/toolkit/src/lib/react-query-service/knowledge/index.ts @@ -11,3 +11,4 @@ export * from "./useGetFileContent"; export * from "./useListChunks"; export * from "./useUpdateChunk"; export * from "./useGetChunkContent"; +export * from "./useGetAllChunks"; diff --git a/packages/toolkit/src/lib/react-query-service/knowledge/useGetAllChunks.ts b/packages/toolkit/src/lib/react-query-service/knowledge/useGetAllChunks.ts new file mode 100644 index 0000000000..34f5f2b6e7 --- /dev/null +++ b/packages/toolkit/src/lib/react-query-service/knowledge/useGetAllChunks.ts @@ -0,0 +1,25 @@ +import { createInstillAxiosClient } from "../../../lib"; + +export async function useGetAllChunks( + accessToken: string, + ownerName: string, + kbId: string, + fileUid: string, +) { + if (!accessToken) { + throw new Error("accessToken not provided"); + } + const client = createInstillAxiosClient(accessToken, true); + try { + const response = await client.get( + `/namespaces/${ownerName}/knowledge-bases/${kbId}/chunks`, + { + params: { fileUid }, + }, + ); + return response.data.chunks; + } catch (error) { + console.error("Error fetching chunks:", error); + throw new Error("Failed to fetch chunks. Please try again later."); + } +} diff --git a/packages/toolkit/src/view/knowledge/components/CreateKnowledgeBaseCard.tsx b/packages/toolkit/src/view/knowledge/components/CreateKnowledgeBaseCard.tsx index efeb71adc2..fd6d66e869 100644 --- a/packages/toolkit/src/view/knowledge/components/CreateKnowledgeBaseCard.tsx +++ b/packages/toolkit/src/view/knowledge/components/CreateKnowledgeBaseCard.tsx @@ -9,8 +9,8 @@ import { Tooltip, } from "@instill-ai/design-system"; -import { InstillStore, useInstillStore, useShallow } from "../../../lib"; -import { useListChunks } from "../../../lib/react-query-service/knowledge"; +import { InstillStore, useInstillStore, useQueries, useShallow } from "../../../lib"; +import { useGetAllChunks, useListKnowledgeBaseFiles } from "../../../lib/react-query-service/knowledge"; import { KnowledgeBase } from "../../../lib/react-query-service/knowledge/types"; import { EditKnowledgeDialog } from "./EditKnowledgeDialog"; @@ -102,38 +102,54 @@ export const CreateKnowledgeBaseCard = ({ const [isHovered, setIsHovered] = React.useState(false); const [mousePosition, setMousePosition] = React.useState({ x: 0, y: 0 }); - const { accessToken, enabledQuery } = useInstillStore( + const { accessToken, enabledQuery, selectedNamespace } = useInstillStore( useShallow((store: InstillStore) => ({ accessToken: store.accessToken, enabledQuery: store.enabledQuery, + selectedNamespace: store.navigationNamespaceAnchor, })), ); - const { data: chunks, isLoading: isLoadingChunks } = useListChunks({ - kbId: knowledgeBase.kbId, + const { data: files } = useListKnowledgeBaseFiles({ + namespaceId: selectedNamespace || null, + knowledgeBaseId: knowledgeBase.kbId, accessToken: accessToken || null, enabled: enabledQuery && isHovered, - ownerId: knowledgeBase.ownerName, - fileUid: "", }); - const tooltipContent = React.useMemo(() => { - if (!isHovered) return ""; + const chunkQueries = useQueries({ + queries: (files || []).map((file) => ({ + queryKey: ["chunks", knowledgeBase.kbId, file.fileUid], + queryFn: () => useGetAllChunks( + accessToken || "", + knowledgeBase.ownerName, + knowledgeBase.kbId, + file.fileUid + ), + enabled: enabledQuery && isHovered && !!files && !!accessToken, + })), + }); - if (isLoadingChunks) return "Loading..."; + const totalChunks = React.useMemo(() => { + return chunkQueries.reduce((total, query) => { + if (query.data) { + return total + query.data.length; + } + return total; + }, 0); + }, [chunkQueries]); - const textChunks = chunks - ? chunks.filter((chunk: { type: string }) => chunk.type === "TEXT") - : []; + const tooltipContent = React.useMemo(() => { + if (!isHovered) return ""; return ` Converting pipeline ID: ${knowledgeBase.convertingPipelines?.[0] || "N/A"} Splitting pipeline ID: ${knowledgeBase.splittingPipelines?.[0] || "N/A"} Embedding pipeline ID: ${knowledgeBase.embeddingPipelines?.[0] || "N/A"} Files #: ${knowledgeBase.totalFiles || "N/A"} -Text Chunks #: ${textChunks.length} +Text Chunks #: ${totalChunks} `.trim(); - }, [isHovered, isLoadingChunks, chunks, knowledgeBase]); + }, [isHovered, knowledgeBase, totalChunks]); const handleDelete = (e: React.MouseEvent) => { e.stopPropagation();