feat:use variavleinsertion

Onelevenvy · Oct 7, 2024 · 76713e4 · 76713e4
1 parent ab39434
commit 76713e4
Show file tree

Hide file tree

Showing 3 changed files with 134 additions and 72 deletions.
diff --git a/web/src/components/WorkFlow/nodes/LLM/LLMNodeProperties.tsx b/web/src/components/WorkFlow/nodes/LLM/LLMNodeProperties.tsx
@@ -1,5 +1,6 @@
 import ModelSelect from "@/components/Common/ModelProvider";
 import { useModelQuery } from "@/hooks/useModelQuery";
+import { useVariableInsertion } from "@/hooks/graphs/useVariableInsertion";
 import {
   Box,
   Input,
@@ -76,27 +77,24 @@ const LLMNodeProperties: React.FC<LLMNodePropertiesProps> = ({
     setValue("openai_api_base", selectedModel?.provider.base_url);
   };
 
-  const handleSystemPromptChange = useCallback((value: string) => {
-    setSystemPromptInput(value);
-    onNodeDataChange(node.id, "systemMessage", value);
-  }, [node.id, onNodeDataChange]);
+  const handleSystemPromptChange = useCallback(
+    (value: string) => {
+      setSystemPromptInput(value);
+      onNodeDataChange(node.id, "systemMessage", value);
+    },
+    [node.id, onNodeDataChange]
+  );
 
-  const insertVariable = useCallback((variable: string) => {
-    const textarea = textareaRef.current;
-    if (textarea) {
-      const start = textarea.selectionStart;
-      const end = textarea.selectionEnd;
-      const text = textarea.value;
-      const before = text.substring(0, start);
-      const after = text.substring(end);
-      const newValue = before + `{${variable}}` + after;
-      setSystemPromptInput(newValue);
-      onNodeDataChange(node.id, "systemMessage", newValue);
-      setShowVariables(false);
-      textarea.focus();
-      textarea.setSelectionRange(start + variable.length + 2, start + variable.length + 2);
-    }
-  }, [node.id, onNodeDataChange]);
+  const {
+    showVariables: showVariablesHook,
+    setShowVariables: setShowVariablesHook,
+    inputRef: inputRefHook,
+    handleKeyDown: handleKeyDownHook,
+    insertVariable: insertVariableHook,
+  } = useVariableInsertion<HTMLTextAreaElement>({
+    onValueChange: (value) => handleSystemPromptChange(value),
+    availableVariables,
+  });
 
   return (
     <VStack align="stretch" spacing={4}>
@@ -130,25 +128,20 @@ const LLMNodeProperties: React.FC<LLMNodePropertiesProps> = ({
       <Box>
         <Text fontWeight="bold">System Prompt:</Text>
         <Popover
-          isOpen={showVariables}
-          onClose={() => setShowVariables(false)}
+          isOpen={showVariablesHook}
+          onClose={() => setShowVariablesHook(false)}
           placement="bottom-start"
         >
           <PopoverTrigger>
             <Textarea
-              ref={textareaRef}
+              ref={inputRefHook}
               value={systemPromptInput}
               onChange={(e) => handleSystemPromptChange(e.target.value)}
-              onKeyDown={(e) => {
-                if (e.key === '/') {
-                  e.preventDefault();
-                  setShowVariables(true);
-                }
-              }}
+              onKeyDown={handleKeyDownHook}
               placeholder="Write your prompt here. Use '/' to insert variables."
               style={{
-                whiteSpace: 'pre-wrap',
-                minHeight: '100px',
+                whiteSpace: "pre-wrap",
+                minHeight: "100px",
               }}
             />
           </PopoverTrigger>
@@ -157,7 +150,9 @@ const LLMNodeProperties: React.FC<LLMNodePropertiesProps> = ({
               {availableVariables.map((v) => (
                 <Button
                   key={`${v.nodeId}.${v.variableName}`}
-                  onClick={() => insertVariable(`${v.nodeId}.${v.variableName}`)}
+                  onClick={() =>
+                    insertVariableHook(`${v.nodeId}.${v.variableName}`)
+                  }
                   size="sm"
                 >
                   {v.nodeId}.{v.variableName}
@@ -171,4 +166,4 @@ const LLMNodeProperties: React.FC<LLMNodePropertiesProps> = ({
   );
 };
 
-export default LLMNodeProperties;
+export default LLMNodeProperties;
diff --git a/web/src/components/WorkFlow/nodes/Plugin/PluginNodeProperties.tsx b/web/src/components/WorkFlow/nodes/Plugin/PluginNodeProperties.tsx
@@ -1,9 +1,19 @@
-import { Box, Text, VStack, Button, Popover, PopoverTrigger, PopoverContent, Input } from "@chakra-ui/react";
+import {
+  Box,
+  Text,
+  VStack,
+  Button,
+  Popover,
+  PopoverTrigger,
+  PopoverContent,
+  Input,
+} from "@chakra-ui/react";
 import type React from "react";
 import { ToolsService } from "@/client/services/ToolsService";
 import { useSkillsQuery } from "@/hooks/useSkillsQuery";
-import { useState, useRef, useCallback } from "react";
-import { VariableReference } from '../../variableSystem';
+import { useState, useCallback } from "react";
+import { VariableReference } from "../../variableSystem";
+import { useVariableInsertion } from "@/hooks/graphs/useVariableInsertion";
 
 interface PluginNodePropertiesProps {
   node: any;
@@ -21,8 +31,31 @@ const PluginNodeProperties: React.FC<PluginNodePropertiesProps> = ({
     (skill) => skill.display_name === node.data.toolName
   );
   const [loading, setLoading] = useState(false);
-  const [showVariables, setShowVariables] = useState<{[key: string]: boolean}>({});
-  const inputRefs = useRef<{[key: string]: HTMLInputElement}>({});
+
+  const handleInputChange = useCallback(
+    (key: string, value: string) => {
+      onNodeDataChange(node.id, "args", {
+        ...node.data.args,
+        [key]: value,
+      });
+    },
+    [node.id, node.data.args, onNodeDataChange]
+  );
+
+  // 为每个可能的输入参数创建一个 useVariableInsertion hook
+  const variableInsertionHooks: {
+    [key: string]: ReturnType<typeof useVariableInsertion<HTMLInputElement>>;
+  } = {};
+
+  if (tool?.input_parameters) {
+    Object.keys(tool.input_parameters).forEach((key) => {
+      // eslint-disable-next-line react-hooks/rules-of-hooks
+      variableInsertionHooks[key] = useVariableInsertion<HTMLInputElement>({
+        onValueChange: (value) => handleInputChange(key, value),
+        availableVariables,
+      });
+    });
+  }
 
   const handleInvoke = async () => {
     setLoading(true);
@@ -39,29 +72,6 @@ const PluginNodeProperties: React.FC<PluginNodePropertiesProps> = ({
     }
   };
 
-  const handleInputChange = useCallback((key: string, value: string) => {
-    onNodeDataChange(node.id, "args", {
-      ...node.data.args,
-      [key]: value,
-    });
-  }, [node.id, node.data.args, onNodeDataChange]);
-
-  const insertVariable = useCallback((key: string, variable: string) => {
-    const input = inputRefs.current[key];
-    if (input) {
-      const start = input.selectionStart || 0;
-      const end = input.selectionEnd || 0;
-      const text = input.value;
-      const before = text.substring(0, start);
-      const after = text.substring(end);
-      const newValue = before + `{${variable}}` + after;
-      handleInputChange(key, newValue);
-      input.focus();
-      input.setSelectionRange(start + variable.length + 2, start + variable.length + 2);
-    }
-    setShowVariables(prev => ({...prev, [key]: false}));
-  }, [handleInputChange]);
-
   return (
     <VStack align="stretch" spacing={4}>
       <Box>
@@ -72,21 +82,18 @@ const PluginNodeProperties: React.FC<PluginNodePropertiesProps> = ({
           <Box key={key}>
             <Text fontWeight="bold">{key}:</Text>
             <Popover
-              isOpen={showVariables[key]}
-              onClose={() => setShowVariables(prev => ({...prev, [key]: false}))}
+              isOpen={variableInsertionHooks[key]?.showVariables}
+              onClose={() =>
+                variableInsertionHooks[key]?.setShowVariables(false)
+              }
               placement="bottom-start"
             >
               <PopoverTrigger>
                 <Input
-                  ref={el => {if (el) inputRefs.current[key] = el;}}
+                  ref={variableInsertionHooks[key]?.inputRef}
                   value={node.data.args[key] || ""}
                   onChange={(e) => handleInputChange(key, e.target.value)}
-                  onKeyDown={(e) => {
-                    if (e.key === '/') {
-                      e.preventDefault();
-                      setShowVariables(prev => ({...prev, [key]: true}));
-                    }
-                  }}
+                  onKeyDown={variableInsertionHooks[key]?.handleKeyDown}
                   placeholder={`Enter ${key}. Use '/' to insert variables.`}
                 />
               </PopoverTrigger>
@@ -95,7 +102,11 @@ const PluginNodeProperties: React.FC<PluginNodePropertiesProps> = ({
                   {availableVariables.map((v) => (
                     <Button
                       key={`${v.nodeId}.${v.variableName}`}
-                      onClick={() => insertVariable(key, `${v.nodeId}.${v.variableName}`)}
+                      onClick={() =>
+                        variableInsertionHooks[key]?.insertVariable(
+                          `${v.nodeId}.${v.variableName}`
+                        )
+                      }
                       size="sm"
                     >
                       {v.nodeId}.{v.variableName}
@@ -113,4 +124,4 @@ const PluginNodeProperties: React.FC<PluginNodePropertiesProps> = ({
   );
 };
 
-export default PluginNodeProperties;
+export default PluginNodeProperties;
diff --git a/web/src/hooks/graphs/useVariableInsertion.ts b/web/src/hooks/graphs/useVariableInsertion.ts
@@ -0,0 +1,56 @@
+import { useState, useRef, useCallback } from "react";
+import { VariableReference } from "../../components/WorkFlow/variableSystem";
+
+interface UseVariableInsertionProps<
+  T extends HTMLInputElement | HTMLTextAreaElement,
+> {
+  onValueChange: (value: string) => void;
+  availableVariables: VariableReference[];
+}
+
+export const useVariableInsertion = <
+  T extends HTMLInputElement | HTMLTextAreaElement,
+>({
+  onValueChange,
+  availableVariables,
+}: UseVariableInsertionProps<T>) => {
+  const [showVariables, setShowVariables] = useState(false);
+  const inputRef = useRef<T>(null);
+
+  const handleKeyDown = useCallback((e: React.KeyboardEvent) => {
+    if (e.key === "/") {
+      e.preventDefault();
+      setShowVariables(true);
+    }
+  }, []);
+
+  const insertVariable = useCallback(
+    (variable: string) => {
+      const input = inputRef.current;
+      if (input) {
+        const start = input.selectionStart || 0;
+        const end = input.selectionEnd || 0;
+        const text = input.value;
+        const before = text.substring(0, start);
+        const after = text.substring(end);
+        const newValue = before + `{${variable}}` + after;
+        onValueChange(newValue);
+        setShowVariables(false);
+        input.focus();
+        input.setSelectionRange(
+          start + variable.length + 2,
+          start + variable.length + 2
+        );
+      }
+    },
+    [onValueChange]
+  );
+
+  return {
+    showVariables,
+    setShowVariables,
+    inputRef,
+    handleKeyDown,
+    insertVariable,
+  };
+};