From 91b3aec292f40eb282f5911d3c8a5268551cbe50 Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Thu, 7 Nov 2024 15:24:02 +0800
Subject: [PATCH] feat: add use tools

---
 .../workflow/block-selector/tool-picker.tsx   | 39 ++++---------
 web/context/query-client.tsx                  |  2 +-
 web/service/use-tools.ts                      | 55 +++++++++++++++++++
 3 files changed, 67 insertions(+), 29 deletions(-)
 create mode 100644 web/service/use-tools.ts

diff --git a/web/app/components/workflow/block-selector/tool-picker.tsx b/web/app/components/workflow/block-selector/tool-picker.tsx
index 95e2685943..3f1ffe4630 100644
--- a/web/app/components/workflow/block-selector/tool-picker.tsx
+++ b/web/app/components/workflow/block-selector/tool-picker.tsx
@@ -1,7 +1,7 @@
 'use client'
 import type { FC } from 'react'
 import React from 'react'
-import { useEffect, useState } from 'react'
+import { useState } from 'react'
 import {
   PortalToFollowElem,
   PortalToFollowElemContent,
@@ -13,12 +13,7 @@ import type {
 } from '@floating-ui/react'
 import AllTools from '@/app/components/workflow/block-selector/all-tools'
 import type { ToolDefaultValue } from './types'
-import {
-  fetchAllBuiltInTools,
-  fetchAllCustomTools,
-  fetchAllWorkflowTools,
-} from '@/service/tools'
-import type { BlockEnum, ToolWithProvider } from '@/app/components/workflow/types'
+import type { BlockEnum } from '@/app/components/workflow/types'
 import SearchBox from '@/app/components/plugins/marketplace/search-box'
 import { useTranslation } from 'react-i18next'
 import { useBoolean } from 'ahooks'
@@ -28,6 +23,7 @@ import {
 } from '@/service/tools'
 import type { CustomCollectionBackend } from '@/app/components/tools/types'
 import Toast from '@/app/components/base/toast'
+import { useAllBuiltInTools, useAllCustomTools, useAllWorkflowTools, useInvalidateAllCustomTools } from '@/service/use-tools'
 
 type Props = {
   disabled: boolean
@@ -53,25 +49,12 @@ const ToolPicker: FC<Props> = ({
   const { t } = useTranslation()
   const [searchText, setSearchText] = useState('')
 
-  const [buildInTools, setBuildInTools] = useState<ToolWithProvider[]>([])
-  const [customTools, setCustomTools] = useState<ToolWithProvider[]>([])
-  const [workflowTools, setWorkflowTools] = useState<ToolWithProvider[]>([])
+  const { data: buildInTools } = useAllBuiltInTools()
+  const { data: customTools } = useAllCustomTools()
+  const { invalidate: invalidateCustomTools } = useInvalidateAllCustomTools()
+  const { data: workflowTools } = useAllWorkflowTools()
 
-  useEffect(() => {
-    (async () => {
-      const buildInTools = await fetchAllBuiltInTools()
-      const customTools = await fetchAllCustomTools()
-      const workflowTools = await fetchAllWorkflowTools()
-      setBuildInTools(buildInTools)
-      setCustomTools(customTools)
-      setWorkflowTools(workflowTools)
-    })()
-  }, [])
-
-  const handleAddedCustomTool = async () => {
-    const customTools = await fetchAllCustomTools()
-    setCustomTools(customTools)
-  }
+  const handleAddedCustomTool = invalidateCustomTools
 
   const handleTriggerClick = () => {
     if (disabled) return
@@ -138,9 +121,9 @@ const ToolPicker: FC<Props> = ({
             className='mt-1'
             searchText={searchText}
             onSelect={handleSelect}
-            buildInTools={buildInTools}
-            customTools={customTools}
-            workflowTools={workflowTools}
+            buildInTools={buildInTools || []}
+            customTools={customTools || []}
+            workflowTools={workflowTools || []}
             supportAddCustomTool={supportAddCustomTool}
             onAddedCustomTool={handleAddedCustomTool}
             onShowAddCustomCollectionModal={showEditCustomCollectionModal}
diff --git a/web/context/query-client.tsx b/web/context/query-client.tsx
index 4cb66eb826..1adb8af653 100644
--- a/web/context/query-client.tsx
+++ b/web/context/query-client.tsx
@@ -6,7 +6,7 @@ import { ReactQueryDevtools } from '@tanstack/react-query-devtools'
 
 const client = new QueryClient()
 
-export const TanstackQueryIniter: FC <PropsWithChildren> = (props) => {
+export const TanstackQueryIniter: FC<PropsWithChildren> = (props) => {
   const { children } = props
   return <QueryClientProvider client={client}>
     {children}
diff --git a/web/service/use-tools.ts b/web/service/use-tools.ts
new file mode 100644
index 0000000000..05e18051f6
--- /dev/null
+++ b/web/service/use-tools.ts
@@ -0,0 +1,55 @@
+import { get } from './base'
+import type {
+  Tool,
+} from '@/app/components/tools/types'
+import type { ToolWithProvider } from '@/app/components/workflow/types'
+import {
+  useQueryClient,
+} from '@tanstack/react-query'
+
+import {
+  useQuery,
+} from '@tanstack/react-query'
+
+const NAME_SPACE = 'tools'
+
+export const useAllBuiltInTools = () => {
+  return useQuery<ToolWithProvider[]>({
+    queryKey: [NAME_SPACE, 'builtIn'],
+    queryFn: () => get<ToolWithProvider[]>('/workspaces/current/tools/builtin'),
+  })
+}
+
+const useAllCustomToolsKey = [NAME_SPACE, 'customTools']
+export const useAllCustomTools = () => {
+  return useQuery<ToolWithProvider[]>({
+    queryKey: useAllCustomToolsKey,
+    queryFn: () => get<ToolWithProvider[]>('/workspaces/current/tools/api'),
+  })
+}
+
+export const useInvalidateAllCustomTools = () => {
+  const queryClient = useQueryClient()
+  return {
+    invalidate: () => {
+      queryClient.invalidateQueries(
+        {
+          queryKey: useAllCustomToolsKey,
+        })
+    },
+  }
+}
+
+export const useAllWorkflowTools = () => {
+  return useQuery<ToolWithProvider[]>({
+    queryKey: [NAME_SPACE, 'workflowTools'],
+    queryFn: () => get<ToolWithProvider[]>('/workspaces/current/tools/workflow'),
+  })
+}
+
+export const useBuiltInTools = (collectionName: string) => {
+  return useQuery({
+    queryKey: [NAME_SPACE, 'builtIn', collectionName],
+    queryFn: () => get<Tool[]>(`/workspaces/current/tool-provider/builtin/${collectionName}/tools`),
+  })
+}