Workflow implementation for data schema suggestions (#1627)

2025-01-23 23:36:13 +08:00
parent b42a1900e2
commit 99adf68444
5 changed files with 185 additions and 160 deletions
--- a/skyvern-frontend/src/components/DataSchemaInputGroup/WorkflowDataSchemaInputGroup.tsx
+++ 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 };