Dorod-Sky/skyvern-frontend/src/components/DataSchemaInputGroup/WorkflowDataSchemaInputGroup.tsx

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 };