From 99adf6844468de94506c65db79ee14c038559807 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Thu, 23 Jan 2025 23:36:13 +0800
Subject: [PATCH] Workflow implementation for data schema suggestions (#1627)

---
 .../WorkflowDataSchemaInputGroup.tsx          | 132 ++++++++++++++++++
 .../nodes/ExtractionNode/ExtractionNode.tsx   |  52 ++-----
 .../nodes/PDFParserNode/PDFParserNode.tsx     |  49 ++-----
 .../editor/nodes/TaskNode/TaskNode.tsx        |  49 ++-----
 .../nodes/TextPromptNode/TextPromptNode.tsx   |  63 +++------
 5 files changed, 185 insertions(+), 160 deletions(-)
 create mode 100644 skyvern-frontend/src/components/DataSchemaInputGroup/WorkflowDataSchemaInputGroup.tsx

diff --git a/skyvern-frontend/src/components/DataSchemaInputGroup/WorkflowDataSchemaInputGroup.tsx b/skyvern-frontend/src/components/DataSchemaInputGroup/WorkflowDataSchemaInputGroup.tsx
new file mode 100644
index 00000000..29a770e8
--- /dev/null
+++ b/skyvern-frontend/src/components/DataSchemaInputGroup/WorkflowDataSchemaInputGroup.tsx
@@ -0,0 +1,132 @@
+import { HelpTooltip } from "@/components/HelpTooltip";
+import { Checkbox } from "@/components/ui/checkbox";
+import { Label } from "@/components/ui/label";
+import {
+  Cross2Icon,
+  MagicWandIcon,
+  PaperPlaneIcon,
+  ReloadIcon,
+} from "@radix-ui/react-icons";
+import { useMutation } from "@tanstack/react-query";
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
+import { getClient } from "@/api/AxiosClient";
+import { CodeEditor } from "@/routes/workflows/components/CodeEditor";
+import { helpTooltips } from "@/routes/workflows/editor/helpContent";
+import { useState } from "react";
+import { AutoResizingTextarea } from "../AutoResizingTextarea/AutoResizingTextarea";
+import { Button } from "../ui/button";
+import { AxiosError } from "axios";
+import { toast } from "../ui/use-toast";
+
+type Props = {
+  value: string;
+  onChange: (value: string) => void;
+  suggestionContext: Record<string, unknown>;
+  exampleValue: Record<string, unknown>;
+};
+
+function WorkflowDataSchemaInputGroup({
+  value,
+  onChange,
+  suggestionContext,
+  exampleValue,
+}: Props) {
+  const credentialGetter = useCredentialGetter();
+  const [generateWithAIActive, setGenerateWithAIActive] = useState(false);
+  const [generateWithAIPrompt, setGenerateWithAIPrompt] = useState("");
+
+  const getDataSchemaSuggestionMutation = useMutation({
+    mutationFn: async () => {
+      const client = await getClient(credentialGetter);
+      return client.post<{ output: Record<string, unknown> }>(
+        "/suggest/data_schema",
+        {
+          input: generateWithAIPrompt,
+          context: suggestionContext,
+        },
+      );
+    },
+    onSuccess: (response) => {
+      onChange(JSON.stringify(response.data.output, null, 2));
+    },
+    onError: (error: AxiosError) => {
+      toast({
+        variant: "destructive",
+        title: "Could not generate the data schema",
+        description:
+          error.message ?? "There was an error generating data schema",
+      });
+    },
+  });
+
+  return (
+    <div className="space-y-2">
+      <div className="flex gap-4">
+        <div className="flex gap-2">
+          <Label className="text-xs text-slate-300">Data Schema</Label>
+          <HelpTooltip content={helpTooltips["task"]["dataSchema"]} />
+        </div>
+        <Checkbox
+          checked={value !== "null"}
+          onCheckedChange={(checked) => {
+            onChange(checked ? JSON.stringify(exampleValue, null, 2) : "null");
+          }}
+        />
+      </div>
+      {value !== "null" && (
+        <div className="space-y-2">
+          <CodeEditor
+            language="json"
+            value={value}
+            onChange={onChange}
+            className="nowheel nopan"
+            fontSize={8}
+          />
+          {value !== "null" &&
+            (generateWithAIActive ? (
+              <div className="flex w-full items-center rounded-xl border px-4">
+                <Cross2Icon
+                  className="size-4"
+                  onClick={() => {
+                    setGenerateWithAIActive(false);
+                    setGenerateWithAIPrompt("");
+                  }}
+                />
+                <AutoResizingTextarea
+                  className="min-h-0 resize-none rounded-md border-transparent px-4 py-2 text-xs hover:border-transparent focus-visible:ring-0"
+                  value={generateWithAIPrompt}
+                  onChange={(event) => {
+                    setGenerateWithAIPrompt(event.target.value);
+                  }}
+                  placeholder="Describe how you want your output formatted"
+                />
+                {getDataSchemaSuggestionMutation.isPending ? (
+                  <ReloadIcon className="size-4 animate-spin" />
+                ) : (
+                  <PaperPlaneIcon
+                    className="size-4"
+                    onClick={() => {
+                      getDataSchemaSuggestionMutation.mutate();
+                    }}
+                  />
+                )}
+              </div>
+            ) : (
+              <Button
+                variant="tertiary"
+                size="sm"
+                onClick={() => {
+                  setGenerateWithAIActive(true);
+                }}
+              >
+                <MagicWandIcon className="mr-2 size-4" />
+                Generate with AI
+              </Button>
+            ))}
+        </div>
+      )}
+    </div>
+  );
+}
+
+export { WorkflowDataSchemaInputGroup };
diff --git a/skyvern-frontend/src/routes/workflows/editor/nodes/ExtractionNode/ExtractionNode.tsx b/skyvern-frontend/src/routes/workflows/editor/nodes/ExtractionNode/ExtractionNode.tsx
index c5212c2b..16fbdd18 100644
--- a/skyvern-frontend/src/routes/workflows/editor/nodes/ExtractionNode/ExtractionNode.tsx
+++ b/skyvern-frontend/src/routes/workflows/editor/nodes/ExtractionNode/ExtractionNode.tsx
@@ -6,12 +6,10 @@ import {
   AccordionItem,
   AccordionTrigger,
 } from "@/components/ui/accordion";
-import { Checkbox } from "@/components/ui/checkbox";
 import { Input } from "@/components/ui/input";
 import { Label } from "@/components/ui/label";
 import { Separator } from "@/components/ui/separator";
 import { Switch } from "@/components/ui/switch";
-import { CodeEditor } from "@/routes/workflows/components/CodeEditor";
 import { useDeleteNodeCallback } from "@/routes/workflows/hooks/useDeleteNodeCallback";
 import { useNodeLabelChangeHandler } from "@/routes/workflows/hooks/useLabelChangeHandler";
 import {
@@ -33,6 +31,7 @@ import { helpTooltips, placeholders } from "../../helpContent";
 import { AppNode } from "..";
 import { getAvailableOutputParameterKeys } from "../../workflowEditorUtils";
 import { ParametersMultiSelect } from "../TaskNode/ParametersMultiSelect";
+import { WorkflowDataSchemaInputGroup } from "@/components/DataSchemaInputGroup/WorkflowDataSchemaInputGroup";
 
 function ExtractionNode({ id, data }: NodeProps<ExtractionNode>) {
   const { updateNodeData } = useReactFlow();
@@ -123,44 +122,17 @@ function ExtractionNode({ id, data }: NodeProps<ExtractionNode>) {
             className="nopan text-xs"
           />
         </div>
-        <div className="space-y-2">
-          <div className="flex gap-4">
-            <div className="flex gap-2">
-              <Label className="text-xs text-slate-300">Data Schema</Label>
-              <HelpTooltip content={helpTooltips["extraction"]["dataSchema"]} />
-            </div>
-            <Checkbox
-              checked={inputs.dataSchema !== "null"}
-              onCheckedChange={(checked) => {
-                if (!editable) {
-                  return;
-                }
-                handleChange(
-                  "dataSchema",
-                  checked
-                    ? JSON.stringify(dataSchemaExampleValue, null, 2)
-                    : "null",
-                );
-              }}
-            />
-          </div>
-          {inputs.dataSchema !== "null" && (
-            <div>
-              <CodeEditor
-                language="json"
-                value={inputs.dataSchema}
-                onChange={(value) => {
-                  if (!editable) {
-                    return;
-                  }
-                  handleChange("dataSchema", value);
-                }}
-                className="nowheel nopan"
-                fontSize={8}
-              />
-            </div>
-          )}
-        </div>
+        <WorkflowDataSchemaInputGroup
+          value={inputs.dataSchema}
+          onChange={(value) => {
+            handleChange("dataSchema", value);
+          }}
+          exampleValue={dataSchemaExampleValue}
+          suggestionContext={{
+            data_extraction_goal: inputs.dataExtractionGoal,
+            current_schema: inputs.dataSchema,
+          }}
+        />
         <Separator />
         <Accordion type="single" collapsible>
           <AccordionItem value="advanced" className="border-b-0">
diff --git a/skyvern-frontend/src/routes/workflows/editor/nodes/PDFParserNode/PDFParserNode.tsx b/skyvern-frontend/src/routes/workflows/editor/nodes/PDFParserNode/PDFParserNode.tsx
index 1b0a86e3..eb3fbf8c 100644
--- a/skyvern-frontend/src/routes/workflows/editor/nodes/PDFParserNode/PDFParserNode.tsx
+++ b/skyvern-frontend/src/routes/workflows/editor/nodes/PDFParserNode/PDFParserNode.tsx
@@ -1,8 +1,6 @@
 import { HelpTooltip } from "@/components/HelpTooltip";
-import { Checkbox } from "@/components/ui/checkbox";
 import { Label } from "@/components/ui/label";
 import { WorkflowBlockInput } from "@/components/WorkflowBlockInput";
-import { CodeEditor } from "@/routes/workflows/components/CodeEditor";
 import { useDeleteNodeCallback } from "@/routes/workflows/hooks/useDeleteNodeCallback";
 import { useNodeLabelChangeHandler } from "@/routes/workflows/hooks/useLabelChangeHandler";
 import { WorkflowBlockTypes } from "@/routes/workflows/types/workflowTypes";
@@ -14,6 +12,7 @@ import { NodeActionMenu } from "../NodeActionMenu";
 import { dataSchemaExampleForFileExtraction } from "../types";
 import { WorkflowBlockIcon } from "../WorkflowBlockIcon";
 import { type PDFParserNode } from "./types";
+import { WorkflowDataSchemaInputGroup } from "@/components/DataSchemaInputGroup/WorkflowDataSchemaInputGroup";
 
 function PDFParserNode({ id, data }: NodeProps<PDFParserNode>) {
   const { updateNodeData } = useReactFlow();
@@ -91,44 +90,14 @@ function PDFParserNode({ id, data }: NodeProps<PDFParserNode>) {
               className="nopan text-xs"
             />
           </div>
-          <div className="space-y-2">
-            <div className="flex gap-4">
-              <div className="flex gap-2">
-                <Label className="text-xs text-slate-300">Data Schema</Label>
-                <HelpTooltip
-                  content={helpTooltips["pdfParser"]["jsonSchema"]}
-                />
-              </div>
-              <Checkbox
-                checked={inputs.jsonSchema !== "null"}
-                onCheckedChange={(checked) => {
-                  handleChange(
-                    "jsonSchema",
-                    checked
-                      ? JSON.stringify(
-                          dataSchemaExampleForFileExtraction,
-                          null,
-                          2,
-                        )
-                      : "null",
-                  );
-                }}
-              />
-            </div>
-            {inputs.jsonSchema !== "null" && (
-              <div>
-                <CodeEditor
-                  language="json"
-                  value={inputs.jsonSchema}
-                  onChange={(value) => {
-                    handleChange("jsonSchema", value);
-                  }}
-                  className="nowheel nopan"
-                  fontSize={8}
-                />
-              </div>
-            )}
-          </div>
+          <WorkflowDataSchemaInputGroup
+            exampleValue={dataSchemaExampleForFileExtraction}
+            value={inputs.jsonSchema}
+            onChange={(value) => {
+              handleChange("jsonSchema", value);
+            }}
+            suggestionContext={{}}
+          />
         </div>
       </div>
     </div>
diff --git a/skyvern-frontend/src/routes/workflows/editor/nodes/TaskNode/TaskNode.tsx b/skyvern-frontend/src/routes/workflows/editor/nodes/TaskNode/TaskNode.tsx
index e77ed30d..b928d054 100644
--- a/skyvern-frontend/src/routes/workflows/editor/nodes/TaskNode/TaskNode.tsx
+++ b/skyvern-frontend/src/routes/workflows/editor/nodes/TaskNode/TaskNode.tsx
@@ -34,6 +34,7 @@ import { dataSchemaExampleValue, errorMappingExampleValue } from "../types";
 import { WorkflowBlockIcon } from "../WorkflowBlockIcon";
 import { ParametersMultiSelect } from "./ParametersMultiSelect";
 import type { TaskNode } from "./types";
+import { WorkflowDataSchemaInputGroup } from "@/components/DataSchemaInputGroup/WorkflowDataSchemaInputGroup";
 
 function TaskNode({ id, data }: NodeProps<TaskNode>) {
   const { updateNodeData } = useReactFlow();
@@ -187,42 +188,18 @@ function TaskNode({ id, data }: NodeProps<TaskNode>) {
                     className="nopan text-xs"
                   />
                 </div>
-                <div className="space-y-2">
-                  <div className="flex gap-4">
-                    <div className="flex gap-2">
-                      <Label className="text-xs text-slate-300">
-                        Data Schema
-                      </Label>
-                      <HelpTooltip
-                        content={helpTooltips["task"]["dataSchema"]}
-                      />
-                    </div>
-                    <Checkbox
-                      checked={inputs.dataSchema !== "null"}
-                      onCheckedChange={(checked) => {
-                        handleChange(
-                          "dataSchema",
-                          checked
-                            ? JSON.stringify(dataSchemaExampleValue, null, 2)
-                            : "null",
-                        );
-                      }}
-                    />
-                  </div>
-                  {inputs.dataSchema !== "null" && (
-                    <div>
-                      <CodeEditor
-                        language="json"
-                        value={inputs.dataSchema}
-                        onChange={(value) => {
-                          handleChange("dataSchema", value);
-                        }}
-                        className="nowheel nopan"
-                        fontSize={8}
-                      />
-                    </div>
-                  )}
-                </div>
+                <WorkflowDataSchemaInputGroup
+                  exampleValue={dataSchemaExampleValue}
+                  onChange={(value) => {
+                    handleChange("dataSchema", value);
+                  }}
+                  value={inputs.dataSchema}
+                  suggestionContext={{
+                    data_extraction_goal: inputs.dataExtractionGoal,
+                    current_schema: inputs.dataSchema,
+                    navigation_goal: inputs.navigationGoal,
+                  }}
+                />
               </div>
             </AccordionContent>
           </AccordionItem>
diff --git a/skyvern-frontend/src/routes/workflows/editor/nodes/TextPromptNode/TextPromptNode.tsx b/skyvern-frontend/src/routes/workflows/editor/nodes/TextPromptNode/TextPromptNode.tsx
index ffd4300a..e431d7f4 100644
--- a/skyvern-frontend/src/routes/workflows/editor/nodes/TextPromptNode/TextPromptNode.tsx
+++ b/skyvern-frontend/src/routes/workflows/editor/nodes/TextPromptNode/TextPromptNode.tsx
@@ -1,8 +1,6 @@
 import { HelpTooltip } from "@/components/HelpTooltip";
-import { Checkbox } from "@/components/ui/checkbox";
 import { Label } from "@/components/ui/label";
 import { Separator } from "@/components/ui/separator";
-import { CodeEditor } from "@/routes/workflows/components/CodeEditor";
 import { useDeleteNodeCallback } from "@/routes/workflows/hooks/useDeleteNodeCallback";
 import { useNodeLabelChangeHandler } from "@/routes/workflows/hooks/useLabelChangeHandler";
 import { WorkflowBlockTypes } from "@/routes/workflows/types/workflowTypes";
@@ -24,6 +22,8 @@ import { ParametersMultiSelect } from "../TaskNode/ParametersMultiSelect";
 import { WorkflowBlockIcon } from "../WorkflowBlockIcon";
 import { type TextPromptNode } from "./types";
 import { WorkflowBlockInputTextarea } from "@/components/WorkflowBlockInputTextarea";
+import { WorkflowDataSchemaInputGroup } from "@/components/DataSchemaInputGroup/WorkflowDataSchemaInputGroup";
+import { dataSchemaExampleValue } from "../types";
 
 function TextPromptNode({ id, data }: NodeProps<TextPromptNode>) {
   const { updateNodeData } = useReactFlow();
@@ -43,6 +43,14 @@ function TextPromptNode({ id, data }: NodeProps<TextPromptNode>) {
     initialValue: data.label,
   });
 
+  function handleChange(key: string, value: unknown) {
+    if (!editable) {
+      return;
+    }
+    setInputs({ ...inputs, [key]: value });
+    updateNodeData(id, { [key]: value });
+  }
+
   return (
     <div>
       <Handle
@@ -92,11 +100,7 @@ function TextPromptNode({ id, data }: NodeProps<TextPromptNode>) {
             isFirstInputInNode
             nodeId={id}
             onChange={(value) => {
-              if (!editable) {
-                return;
-              }
-              setInputs({ ...inputs, prompt: value });
-              updateNodeData(id, { prompt: value });
+              handleChange("prompt", value);
             }}
             value={inputs.prompt}
             placeholder="What do you want to generate?"
@@ -113,43 +117,14 @@ function TextPromptNode({ id, data }: NodeProps<TextPromptNode>) {
           />
         </div>
         <Separator />
-        <div className="space-y-2">
-          <div className="flex gap-2">
-            <Label className="text-xs text-slate-300">Data Schema</Label>
-            <Checkbox
-              checked={inputs.jsonSchema !== "null"}
-              onCheckedChange={(checked) => {
-                if (!editable) {
-                  return;
-                }
-                setInputs({
-                  ...inputs,
-                  jsonSchema: checked ? "{}" : "null",
-                });
-                updateNodeData(id, {
-                  jsonSchema: checked ? "{}" : "null",
-                });
-              }}
-            />
-          </div>
-          {inputs.jsonSchema !== "null" && (
-            <div>
-              <CodeEditor
-                language="json"
-                value={inputs.jsonSchema}
-                onChange={(value) => {
-                  if (!editable) {
-                    return;
-                  }
-                  setInputs({ ...inputs, jsonSchema: value });
-                  updateNodeData(id, { jsonSchema: value });
-                }}
-                className="nowheel nopan"
-                fontSize={8}
-              />
-            </div>
-          )}
-        </div>
+        <WorkflowDataSchemaInputGroup
+          exampleValue={dataSchemaExampleValue}
+          value={inputs.jsonSchema}
+          onChange={(value) => {
+            handleChange("jsonSchema", value);
+          }}
+          suggestionContext={{}}
+        />
       </div>
     </div>
   );