From a326d455c878079c7c02c99dc1fc872599bc8e58 Mon Sep 17 00:00:00 2001 From: Ryan Hopper-Lowe Date: Fri, 20 Dec 2024 10:15:58 -0600 Subject: [PATCH] chore: refactor knowledge code in admin ui This should have no effect on logic or UI/UX. it just extracts some subcomponents and abstracts data logic into reusable hooks. Doing this in preparation for thread level knowledge implementations --- .../components/knowledge/AddSourceModal.tsx | 42 +- .../knowledge/AgentKnowledgePanel.tsx | 550 +++--------------- .../knowledge/KnowledgeSourceDetail.tsx | 4 +- .../components/AddKnowledgeButton.tsx | 99 ++++ .../components/KnowledgeFileItem.tsx | 100 ++++ .../components/KnowledgeSourceItem.tsx | 97 +++ .../knowledge/hooks/useKnowledgeFiles.ts | 83 +++ .../knowledge/hooks/useKnowledgeSources.ts | 128 ++++ 8 files changed, 611 insertions(+), 492 deletions(-) create mode 100644 ui/admin/app/components/knowledge/components/AddKnowledgeButton.tsx create mode 100644 ui/admin/app/components/knowledge/components/KnowledgeFileItem.tsx create mode 100644 ui/admin/app/components/knowledge/components/KnowledgeSourceItem.tsx create mode 100644 ui/admin/app/components/knowledge/hooks/useKnowledgeFiles.ts create mode 100644 ui/admin/app/components/knowledge/hooks/useKnowledgeSources.ts diff --git a/ui/admin/app/components/knowledge/AddSourceModal.tsx b/ui/admin/app/components/knowledge/AddSourceModal.tsx index 38187326..ede13a74 100644 --- a/ui/admin/app/components/knowledge/AddSourceModal.tsx +++ b/ui/admin/app/components/knowledge/AddSourceModal.tsx @@ -2,7 +2,6 @@ import { FC, useState } from "react"; import { KNOWLEDGE_TOOL } from "~/lib/model/agents"; import { KnowledgeSourceType } from "~/lib/model/knowledge"; -import { KnowledgeService } from "~/lib/service/api/knowledgeService"; import KnowledgeSourceAvatar from "~/components/knowledge/KnowledgeSourceAvatar"; import { Button } from "~/components/ui/button"; @@ -11,58 +10,39 @@ import { Input } from "~/components/ui/input"; import { Label } from "~/components/ui/label"; interface AddSourceModalProps { - agentId: string; sourceType: KnowledgeSourceType; startPolling: () => void; isOpen: boolean; onOpenChange: (open: boolean) => void; - onSave: (knowledgeSourceId: string) => void; addTool: (tool: string) => void; + onAddWebsite: (website: string) => void; + onAddOneDrive: (link: string) => void; } -const AddSourceModal: FC = ({ - agentId, +export const AddSourceModal: FC = ({ sourceType, startPolling, isOpen, onOpenChange, - onSave, addTool, + onAddWebsite, + onAddOneDrive, }) => { const [newWebsite, setNewWebsite] = useState(""); const [newLink, setNewLink] = useState(""); const handleAddWebsite = async () => { if (newWebsite) { - const trimmedWebsite = newWebsite.trim(); - const formattedWebsite = - trimmedWebsite.startsWith("http://") || - trimmedWebsite.startsWith("https://") - ? trimmedWebsite - : `https://${trimmedWebsite}`; - - const res = await KnowledgeService.createKnowledgeSource(agentId, { - websiteCrawlingConfig: { - urls: [formattedWebsite], - }, - }); - onSave(res.id); - startPolling(); + onAddWebsite(newWebsite); setNewWebsite(""); - onOpenChange(false); } }; const handleAddOneDrive = async () => { - const res = await KnowledgeService.createKnowledgeSource(agentId, { - onedriveConfig: { - sharedLinks: [newLink.trim()], - }, - }); - onSave(res.id); - setNewLink(""); - startPolling(); - onOpenChange(false); + if (newLink) { + onAddOneDrive(newLink); + setNewLink(""); + } }; const handleAdd = async () => { @@ -163,5 +143,3 @@ const AddSourceModal: FC = ({ ); }; - -export default AddSourceModal; diff --git a/ui/admin/app/components/knowledge/AgentKnowledgePanel.tsx b/ui/admin/app/components/knowledge/AgentKnowledgePanel.tsx index c346c2fe..0216cb04 100644 --- a/ui/admin/app/components/knowledge/AgentKnowledgePanel.tsx +++ b/ui/admin/app/components/knowledge/AgentKnowledgePanel.tsx @@ -1,60 +1,29 @@ -import { Avatar } from "@radix-ui/react-avatar"; -import { - Edit, - EyeIcon, - FileIcon, - GlobeIcon, - PlusIcon, - RefreshCcw, - RotateCcwIcon, - Trash, - UploadIcon, -} from "lucide-react"; -import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { useCallback, useRef, useState } from "react"; import { $path } from "safe-routes"; -import useSWR, { SWRResponse } from "swr"; +import useSWR from "swr"; import { Agent, KNOWLEDGE_TOOL } from "~/lib/model/agents"; import { - KnowledgeFile, - KnowledgeFileState, - KnowledgeSource, - KnowledgeSourceStatus, KnowledgeSourceType, - getKnowledgeSourceDisplayName, getKnowledgeSourceType, } from "~/lib/model/knowledge"; import { ModelAlias } from "~/lib/model/models"; import { DefaultModelAliasApiService } from "~/lib/service/api/defaultModelAliasApiService"; -import { KnowledgeService } from "~/lib/service/api/knowledgeService"; -import { assetUrl } from "~/lib/utils"; import { TypographyP } from "~/components/Typography"; import { ErrorDialog } from "~/components/composed/ErrorDialog"; import { WarningAlert } from "~/components/composed/WarningAlert"; -import AddSourceModal from "~/components/knowledge/AddSourceModal"; -import FileStatusIcon from "~/components/knowledge/FileStatusIcon"; -import RemoteFileAvatar from "~/components/knowledge/KnowledgeSourceAvatar"; -import KnowledgeSourceDetail from "~/components/knowledge/KnowledgeSourceDetail"; -import { LoadingSpinner } from "~/components/ui/LoadingSpinner"; -import { Button } from "~/components/ui/button"; -import { - DropdownMenu, - DropdownMenuContent, - DropdownMenuItem, - DropdownMenuTrigger, -} from "~/components/ui/dropdown-menu"; +import { AddSourceModal } from "~/components/knowledge/AddSourceModal"; +import { KnowledgeSourceDetail } from "~/components/knowledge/KnowledgeSourceDetail"; +import { AddKnowledgeButton } from "~/components/knowledge/components/AddKnowledgeButton"; +import { KnowledgeFileItem } from "~/components/knowledge/components/KnowledgeFileItem"; +import { KnowledgeSourceItem } from "~/components/knowledge/components/KnowledgeSourceItem"; +import { useKnowledgeFiles } from "~/components/knowledge/hooks/useKnowledgeFiles"; +import { useKnowledgeSources } from "~/components/knowledge/hooks/useKnowledgeSources"; import { Input } from "~/components/ui/input"; import { Label } from "~/components/ui/label"; import { Link } from "~/components/ui/link"; import { AutosizeTextarea } from "~/components/ui/textarea"; -import { - Tooltip, - TooltipContent, - TooltipProvider, - TooltipTrigger, -} from "~/components/ui/tooltip"; -import { useAsync } from "~/hooks/useAsync"; import { useMultiAsync } from "~/hooks/useMultiAsync"; type AgentKnowledgePanelProps = { @@ -71,12 +40,7 @@ export default function AgentKnowledgePanel({ addTool, }: AgentKnowledgePanelProps) { const fileInputRef = useRef(null); - const [blockPollingLocalFiles, setBlockPollingLocalFiles] = useState(false); - const [blockPollingSources, setBlockPollingSources] = useState(false); const [isAddSourceModalOpen, setIsAddSourceModalOpen] = useState(false); - const [knowledgeDescription, setKnowledgeDescription] = useState( - agent.knowledgeDescription - ); const [sourceType, setSourceType] = useState( KnowledgeSourceType.Website ); @@ -85,7 +49,6 @@ export default function AgentKnowledgePanel({ >(undefined); const [isEditKnowledgeSourceModalOpen, setIsEditKnowledgeSourceModalOpen] = useState(false); - const [errorDialogError, setErrorDialogError] = useState(""); const { data: defaultAliases } = useSWR( @@ -93,116 +56,31 @@ export default function AgentKnowledgePanel({ DefaultModelAliasApiService.getAliases ); - const getLocalFiles: SWRResponse = useSWR( - KnowledgeService.getLocalKnowledgeFilesForAgent.key(agentId), - ({ agentId }) => - KnowledgeService.getLocalKnowledgeFilesForAgent(agentId), - { - revalidateOnFocus: false, - refreshInterval: blockPollingLocalFiles ? undefined : 5000, - } - ); - const localFiles = useMemo( - () => - getLocalFiles.data - ?.sort((a, b) => a.fileName.localeCompare(b.fileName)) - .map( - (item) => - ({ - ...item, - }) as KnowledgeFile - ) - .filter((item) => !item.deleted) || [], - [getLocalFiles.data] - ); - - const getKnowledgeSources = useSWR( - KnowledgeService.getKnowledgeSourcesForAgent.key(agentId), - ({ agentId }) => KnowledgeService.getKnowledgeSourcesForAgent(agentId), - { - revalidateOnFocus: false, - refreshInterval: blockPollingSources ? undefined : 5000, - } - ); - const knowledgeSources = useMemo( - () => - getKnowledgeSources.data?.filter((source) => !source.deleted) || [], - [getKnowledgeSources.data] + const { localFiles, addKnowledgeFile, deleteKnowledgeFile, reingestFile } = + useKnowledgeFiles(agentId); + + const { + knowledgeSources, + syncKnowledgeSource, + deleteKnowledgeSource, + updateKnowledgeSource, + addWebsite, + addOneDrive, + addNotion, + } = useKnowledgeSources(agentId); + + const selectedKnowledgeSource = knowledgeSources.find( + (source) => source.id === selectedKnowledgeSourceId ); - const handleRemoteKnowledgeSourceSync = async (id: string) => { - const syncedSource = await KnowledgeService.resyncKnowledgeSource( - agentId, - id - ); - getKnowledgeSources.mutate((prev) => - prev?.map((source) => - source.id === syncedSource.id ? syncedSource : source - ) - ); - }; - - const handleDeleteKnowledgeSource = async (id: string) => { - await KnowledgeService.deleteKnowledgeSource(agentId, id); - getKnowledgeSources.mutate( - (prev) => prev?.filter((source) => source.id !== id), - false - ); - }; - - useEffect(() => { - if ( - localFiles.every( - (file) => - file.state === KnowledgeFileState.Ingested || - file.state === KnowledgeFileState.Error - ) - ) { - setBlockPollingLocalFiles(true); - } else { - setBlockPollingLocalFiles(false); - } - }, [localFiles]); - - useEffect(() => { - if ( - knowledgeSources.length === 0 || - knowledgeSources.every( - (source) => - source.state === KnowledgeSourceStatus.Synced || - source.state === KnowledgeSourceStatus.Error - ) - ) { - setBlockPollingSources(true); - } else { - setBlockPollingSources(false); - } - }, [knowledgeSources]); - - const onSaveKnowledgeSource = (updatedSource: KnowledgeSource) => { - getKnowledgeSources.mutate((prev) => - prev?.map((source) => - source.id === updatedSource.id ? updatedSource : source - ) - ); - }; - - //Local file upload const handleAddKnowledge = useCallback( async (_index: number, file: File) => { - const addedFile = await KnowledgeService.addKnowledgeFilesToAgent( - agentId, - file - ); - getLocalFiles.mutate((prev) => - prev ? [...prev, addedFile] : [addedFile] - ); + const addedFile = await addKnowledgeFile(file); return addedFile; }, - [agentId, getLocalFiles] + [addKnowledgeFile] ); - // use multi async to handle uploading multiple files at once const uploadKnowledge = useMultiAsync(handleAddKnowledge); const startUpload = (files: FileList) => { @@ -217,23 +95,31 @@ export default function AgentKnowledgePanel({ if (fileInputRef.current) fileInputRef.current.value = ""; }; - const deleteKnowledge = useAsync(async (item: KnowledgeFile) => { - await KnowledgeService.deleteKnowledgeFileFromAgent( - agentId, - item.fileName - ); - getLocalFiles.mutate((prev) => prev?.filter((f) => f.id !== item.id)); - }); - - const selectedKnowledgeSource = useMemo(() => { - return knowledgeSources.find( - (source) => source.id === selectedKnowledgeSourceId - ); - }, [knowledgeSources, selectedKnowledgeSourceId]); - const hasDefaultTextEmbedding = defaultAliases?.some( (alias) => alias.alias === ModelAlias.TextEmbedding && !!alias.model ); + + const handleSave = (knowledgeSourceId: string): void => { + addTool(KNOWLEDGE_TOOL); + setSelectedKnowledgeSourceId(knowledgeSourceId); + setIsEditKnowledgeSourceModalOpen(true); + }; + + const handleAddWebsite = async (website: string) => { + const res = await addWebsite(website); + handleSave(res.id); + }; + + const handleAddOneDrive = async (link: string) => { + const res = await addOneDrive(link); + handleSave(res.id); + }; + + const handleAddNotion = async () => { + const res = await addNotion(); + handleSave(res.id); + }; + return (
{!hasDefaultTextEmbedding && ( @@ -254,320 +140,70 @@ export default function AgentKnowledgePanel({ ) => { - setKnowledgeDescription(e.target.value); + onChange={(e) => updateAgent({ ...agent, knowledgeDescription: e.target.value, - }); - }} + }) + } className="max-h-[400px]" />
{localFiles.map((file) => ( -
-
- - {file.fileName} -
-
-
- {file.sizeInBytes - ? file.sizeInBytes > 1000000 - ? (file.sizeInBytes / 1000000).toFixed( - 2 - ) + " MB" - : (file.sizeInBytes / 1000).toFixed(2) + - " KB" - : "0 Bytes"} -
-
- {file.state === KnowledgeFileState.Error ? ( -
- - - - - - Reingest - - - - - - - - - View Error - - -
- ) : ( -
- -
- )} -
- - - - - - Delete - - -
-
+ file={file} + onDelete={deleteKnowledgeFile} + onReingest={(file) => reingestFile(file.id!)} + onViewError={setErrorDialogError} + /> ))} + {knowledgeSources.map((source) => ( -
-
- - {getKnowledgeSourceDisplayName(source)} -
-
- - - - - - {source.state === - KnowledgeSourceStatus.Syncing || - source.state === - KnowledgeSourceStatus.Pending - ? (source.status ?? "Syncing...") - : "Sync"} - - - - - - - - Edit - - - - - - - Delete - -
-
+ source={source} + onSync={syncKnowledgeSource} + onEdit={(id) => { + setSelectedKnowledgeSourceId(id); + setIsEditKnowledgeSourceModalOpen(true); + }} + onDelete={deleteKnowledgeSource} + /> ))} -
- - - - - - fileInputRef.current?.click()} - className="cursor-pointer" - > -
- - Local Files -
- - { - setSourceType(KnowledgeSourceType.OneDrive); - setIsAddSourceModalOpen(true); - }} - className="cursor-pointer" - > -
-
-
- - OneDrive logo - -
- OneDrive -
-
- - { - const res = - await KnowledgeService.createKnowledgeSource( - agentId, - { notionConfig: {} } - ); - addTool(KNOWLEDGE_TOOL); - getKnowledgeSources.mutate(); - setSelectedKnowledgeSourceId(res.id); - setIsEditKnowledgeSourceModalOpen(true); - }} - className="cursor-pointer" - disabled={knowledgeSources.some( - (source) => - getKnowledgeSourceType(source) === - KnowledgeSourceType.Notion - )} - > -
- - Notion logo - - Notion -
- - { - setSourceType(KnowledgeSourceType.Website); - setIsAddSourceModalOpen(true); - }} - className="cursor-pointer" - > -
- - Website -
- - - -
+ fileInputRef.current?.click()} + onAddSource={(type) => { + if (type === KnowledgeSourceType.Notion) { + handleAddNotion(); + } else { + setSourceType(type); + setIsAddSourceModalOpen(true); + } + }} + hasExistingNotion={knowledgeSources.some( + (source) => + getKnowledgeSourceType(source) === + KnowledgeSourceType.Notion + )} + />
{ - getKnowledgeSources.mutate(); - }} - onSave={(knowledgeSourceId) => { - setSelectedKnowledgeSourceId(knowledgeSourceId); - setIsEditKnowledgeSourceModalOpen(true); - }} + startPolling={() => {}} addTool={addTool} + onAddWebsite={handleAddWebsite} + onAddOneDrive={handleAddOneDrive} /> - handleRemoteKnowledgeSourceSync( - selectedKnowledgeSourceId - ) + syncKnowledgeSource(selectedKnowledgeSourceId) } onDelete={() => - handleDeleteKnowledgeSource(selectedKnowledgeSourceId) + deleteKnowledgeSource(selectedKnowledgeSourceId) + } + onSave={(source) => + updateKnowledgeSource(source.id, source) } - onSave={onSaveKnowledgeSource} /> )} void; } -const KnowledgeSourceDetail: FC = ({ +export const KnowledgeSourceDetail: FC = ({ agentId, knowledgeSource, isOpen, @@ -603,5 +603,3 @@ const KnowledgeSourceDetail: FC = ({ ); }; - -export default KnowledgeSourceDetail; diff --git a/ui/admin/app/components/knowledge/components/AddKnowledgeButton.tsx b/ui/admin/app/components/knowledge/components/AddKnowledgeButton.tsx new file mode 100644 index 00000000..0f29895c --- /dev/null +++ b/ui/admin/app/components/knowledge/components/AddKnowledgeButton.tsx @@ -0,0 +1,99 @@ +import { Avatar } from "@radix-ui/react-avatar"; +import { GlobeIcon, PlusIcon, UploadIcon } from "lucide-react"; + +import { KnowledgeSourceType } from "~/lib/model/knowledge"; +import { assetUrl } from "~/lib/utils"; + +import { Button } from "~/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "~/components/ui/dropdown-menu"; + +interface AddKnowledgeButtonProps { + disabled?: boolean; + onUploadFiles: () => void; + onAddSource: (sourceType: KnowledgeSourceType) => void; + hasExistingNotion: boolean; +} + +export function AddKnowledgeButton({ + disabled, + onUploadFiles, + onAddSource, + hasExistingNotion, +}: AddKnowledgeButtonProps) { + return ( +
+ + + + + + +
+ + Local Files +
+ + + onAddSource(KnowledgeSourceType.OneDrive) + } + className="cursor-pointer" + > +
+
+
+ + OneDrive logo + +
+ OneDrive +
+
+ + onAddSource(KnowledgeSourceType.Notion)} + className="cursor-pointer" + disabled={hasExistingNotion} + > +
+ + Notion logo + + Notion +
+ + onAddSource(KnowledgeSourceType.Website)} + className="cursor-pointer" + > +
+ + Website +
+ + + +
+ ); +} diff --git a/ui/admin/app/components/knowledge/components/KnowledgeFileItem.tsx b/ui/admin/app/components/knowledge/components/KnowledgeFileItem.tsx new file mode 100644 index 00000000..f9cb5c08 --- /dev/null +++ b/ui/admin/app/components/knowledge/components/KnowledgeFileItem.tsx @@ -0,0 +1,100 @@ +import { EyeIcon, FileIcon, RotateCcwIcon, Trash } from "lucide-react"; + +import { KnowledgeFile, KnowledgeFileState } from "~/lib/model/knowledge"; + +import FileStatusIcon from "~/components/knowledge/FileStatusIcon"; +import { Button } from "~/components/ui/button"; +import { + Tooltip, + TooltipContent, + TooltipProvider, + TooltipTrigger, +} from "~/components/ui/tooltip"; + +interface KnowledgeFileItemProps { + file: KnowledgeFile; + onDelete: (file: KnowledgeFile) => void; + onReingest: (file: KnowledgeFile) => void; + onViewError: (error: string) => void; +} + +export function KnowledgeFileItem({ + file, + onDelete, + onReingest, + onViewError, +}: KnowledgeFileItemProps) { + const formatFileSize = (bytes: number) => { + if (bytes > 1000000) { + return (bytes / 1000000).toFixed(2) + " MB"; + } + return (bytes / 1000).toFixed(2) + " KB"; + }; + + return ( +
+
+ + {file.fileName} +
+
+
+ {file.sizeInBytes + ? formatFileSize(file.sizeInBytes) + : "0 Bytes"} +
+
+ {file.state === KnowledgeFileState.Error ? ( +
+ + + + + Reingest + + + + + + + View Error + +
+ ) : ( +
+ +
+ )} +
+ + + + + + Delete + + +
+
+ ); +} diff --git a/ui/admin/app/components/knowledge/components/KnowledgeSourceItem.tsx b/ui/admin/app/components/knowledge/components/KnowledgeSourceItem.tsx new file mode 100644 index 00000000..fce1cba3 --- /dev/null +++ b/ui/admin/app/components/knowledge/components/KnowledgeSourceItem.tsx @@ -0,0 +1,97 @@ +import { Edit, RefreshCcw, Trash } from "lucide-react"; + +import { + KnowledgeSource, + KnowledgeSourceStatus, + getKnowledgeSourceDisplayName, + getKnowledgeSourceType, +} from "~/lib/model/knowledge"; + +import KnowledgeSourceAvatar from "~/components/knowledge/KnowledgeSourceAvatar"; +import { LoadingSpinner } from "~/components/ui/LoadingSpinner"; +import { Button } from "~/components/ui/button"; +import { + Tooltip, + TooltipContent, + TooltipProvider, + TooltipTrigger, +} from "~/components/ui/tooltip"; + +interface KnowledgeSourceItemProps { + source: KnowledgeSource; + onSync: (sourceId: string) => void; + onEdit: (sourceId: string) => void; + onDelete: (sourceId: string) => void; +} + +export function KnowledgeSourceItem({ + source, + onSync, + onEdit, + onDelete, +}: KnowledgeSourceItemProps) { + const isSyncing = + source.state === KnowledgeSourceStatus.Syncing || + source.state === KnowledgeSourceStatus.Pending; + + return ( +
+
+ + {getKnowledgeSourceDisplayName(source)} +
+
+ + + + + + {isSyncing ? (source.status ?? "Syncing...") : "Sync"} + + + + + + + + Edit + + + + + + + + Delete + + +
+
+ ); +} diff --git a/ui/admin/app/components/knowledge/hooks/useKnowledgeFiles.ts b/ui/admin/app/components/knowledge/hooks/useKnowledgeFiles.ts new file mode 100644 index 00000000..8026d37f --- /dev/null +++ b/ui/admin/app/components/knowledge/hooks/useKnowledgeFiles.ts @@ -0,0 +1,83 @@ +import { useMemo, useState } from "react"; +import useSWR from "swr"; + +import { KnowledgeFile, KnowledgeFileState } from "~/lib/model/knowledge"; +import { KnowledgeService } from "~/lib/service/api/knowledgeService"; + +export function useKnowledgeFiles(agentId: string) { + const [blockPollingLocalFiles, setBlockPollingLocalFiles] = useState(false); + + const { + data: files, + mutate: mutateFiles, + ...rest + } = useSWR( + KnowledgeService.getLocalKnowledgeFilesForAgent.key(agentId), + ({ agentId }) => + KnowledgeService.getLocalKnowledgeFilesForAgent(agentId), + { + revalidateOnFocus: false, + refreshInterval: blockPollingLocalFiles ? undefined : 5000, + } + ); + + const localFiles = useMemo(() => { + return ( + files + ?.sort((a, b) => a.fileName.localeCompare(b.fileName)) + .map((item) => ({ ...item }) as KnowledgeFile) + .filter((item) => !item.deleted) || [] + ); + }, [files]); + + const shouldBlock = useMemo( + () => + localFiles.every( + (file) => + file.state === KnowledgeFileState.Ingested || + file.state === KnowledgeFileState.Error + ), + [localFiles] + ); + + if (shouldBlock !== blockPollingLocalFiles) { + setBlockPollingLocalFiles(shouldBlock); + } + + const addKnowledgeFile = async (file: File) => { + const addedFile = await KnowledgeService.addKnowledgeFilesToAgent( + agentId, + file + ); + mutateFiles((prev) => (prev ? [...prev, addedFile] : [addedFile])); + return addedFile; + }; + + const deleteKnowledgeFile = async (file: KnowledgeFile) => { + await KnowledgeService.deleteKnowledgeFileFromAgent( + agentId, + file.fileName + ); + mutateFiles((prev) => prev?.filter((f) => f.id !== file.id)); + }; + + const reingestFile = async (fileId: string) => { + const reingestedFile = await KnowledgeService.reingestFile( + agentId, + fileId + ); + mutateFiles((prev) => + prev?.map((f) => (f.id === reingestedFile.id ? reingestedFile : f)) + ); + return reingestedFile; + }; + + return { + localFiles, + addKnowledgeFile, + deleteKnowledgeFile, + reingestFile, + mutateFiles, + ...rest, + }; +} diff --git a/ui/admin/app/components/knowledge/hooks/useKnowledgeSources.ts b/ui/admin/app/components/knowledge/hooks/useKnowledgeSources.ts new file mode 100644 index 00000000..fe7f50ad --- /dev/null +++ b/ui/admin/app/components/knowledge/hooks/useKnowledgeSources.ts @@ -0,0 +1,128 @@ +import { useMemo, useState } from "react"; +import useSWR from "swr"; + +import { + KnowledgeSource, + KnowledgeSourceInput, + KnowledgeSourceStatus, +} from "~/lib/model/knowledge"; +import { KnowledgeService } from "~/lib/service/api/knowledgeService"; + +export function useKnowledgeSources(agentId: string) { + const [blockPollingSources, setBlockPollingSources] = useState(false); + const startPolling = () => setBlockPollingSources(false); + + const { + data: sources, + mutate: mutateSources, + ...rest + } = useSWR( + KnowledgeService.getKnowledgeSourcesForAgent.key(agentId), + ({ agentId }) => KnowledgeService.getKnowledgeSourcesForAgent(agentId), + { + revalidateOnFocus: false, + refreshInterval: blockPollingSources ? undefined : 5000, + } + ); + + const knowledgeSources = useMemo(() => { + return sources?.filter((source) => !source.deleted) || []; + }, [sources]); + + const shouldBlockPolling = + knowledgeSources.length === 0 || + knowledgeSources.every( + (source) => + source.state === KnowledgeSourceStatus.Synced || + source.state === KnowledgeSourceStatus.Error + ); + + if (shouldBlockPolling !== blockPollingSources) { + setBlockPollingSources(shouldBlockPolling); + } + + const syncKnowledgeSource = async (sourceId: string) => { + const syncedSource = await KnowledgeService.resyncKnowledgeSource( + agentId, + sourceId + ); + mutateSources((prev) => + prev?.map((source) => + source.id === syncedSource.id ? syncedSource : source + ) + ); + return syncedSource; + }; + + const deleteKnowledgeSource = async (sourceId: string) => { + await KnowledgeService.deleteKnowledgeSource(agentId, sourceId); + mutateSources( + (prev) => prev?.filter((source) => source.id !== sourceId), + false + ); + }; + + const updateKnowledgeSource = async ( + sourceId: string, + updates: Partial + ) => { + const source = knowledgeSources.find((s) => s.id === sourceId); + if (!source) throw new Error("Source not found"); + + const updatedSource = await KnowledgeService.updateKnowledgeSource( + agentId, + sourceId, + { ...source, ...updates } + ); + mutateSources((prev) => + prev?.map((s) => (s.id === updatedSource.id ? updatedSource : s)) + ); + return updatedSource; + }; + + const createKnowledgeSource = async (config: KnowledgeSourceInput) => { + const newSource = await KnowledgeService.createKnowledgeSource( + agentId, + config + ); + mutateSources(); + startPolling(); + return newSource; + }; + + const addWebsite = async (website: string) => { + const trimmedWebsite = website.trim(); + const formattedWebsite = + trimmedWebsite.startsWith("http://") || + trimmedWebsite.startsWith("https://") + ? trimmedWebsite + : `https://${trimmedWebsite}`; + + return await createKnowledgeSource({ + websiteCrawlingConfig: { urls: [formattedWebsite] }, + }); + }; + + const addOneDrive = async (link: string) => { + return await createKnowledgeSource({ + onedriveConfig: { sharedLinks: [link.trim()] }, + }); + }; + + const addNotion = async () => { + return await createKnowledgeSource({ notionConfig: {} }); + }; + + return { + knowledgeSources, + syncKnowledgeSource, + deleteKnowledgeSource, + updateKnowledgeSource, + createKnowledgeSource, + mutateSources, + addWebsite, + addOneDrive, + addNotion, + ...rest, + }; +}