From 047e534194f58a5d2c00c2ade6d4a70c45992c15 Mon Sep 17 00:00:00 2001
From: Kerem Yilmaz <kerem@skyvern.com>
Date: Fri, 23 Aug 2024 23:54:07 +0300
Subject: [PATCH] Task generation UI changes (#721)

Co-authored-by: Muhammed Salih Altun <muhammedsalihaltun@gmail.com>
---
 skyvern-frontend/src/api/types.ts             |   9 ++
 .../src/routes/tasks/create/TaskTemplates.tsx | 112 +++++++++++++++++-
 2 files changed, 115 insertions(+), 6 deletions(-)

diff --git a/skyvern-frontend/src/api/types.ts b/skyvern-frontend/src/api/types.ts
index f50a5735..c35b4356 100644
--- a/skyvern-frontend/src/api/types.ts
+++ b/skyvern-frontend/src/api/types.ts
@@ -244,3 +244,12 @@ export type WorkflowRunStatusApiResponse = {
   recording_url: string | null;
   outputs: Record<string, unknown> | null;
 };
+
+export type TaskGenerationApiResponse = {
+  suggested_title: string | null;
+  url: string | null;
+  navigation_goal: string | null;
+  data_extraction_goal: string | null;
+  navigation_payload: Record<string, unknown> | null;
+  extracted_information_schema: Record<string, unknown> | null;
+};
diff --git a/skyvern-frontend/src/routes/tasks/create/TaskTemplates.tsx b/skyvern-frontend/src/routes/tasks/create/TaskTemplates.tsx
index dbd4e156..aecb88f6 100644
--- a/skyvern-frontend/src/routes/tasks/create/TaskTemplates.tsx
+++ b/skyvern-frontend/src/routes/tasks/create/TaskTemplates.tsx
@@ -17,12 +17,14 @@ import {
   PaperPlaneIcon,
   ReloadIcon,
 } from "@radix-ui/react-icons";
-import { useMutation } from "@tanstack/react-query";
+import { useMutation, useQueryClient } from "@tanstack/react-query";
 import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 import { getClient } from "@/api/AxiosClient";
 import { AxiosError } from "axios";
 import { toast } from "@/components/ui/use-toast";
 import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert";
+import { TaskGenerationApiResponse } from "@/api/types";
+import { stringify as convertToYAML } from "yaml";
 
 const examplePrompts = [
   "What is the top post on hackernews?",
@@ -61,20 +63,113 @@ const templateSamples: {
   },
 };
 
+function createTemplateTaskFromTaskGenerationParameters(
+  values: TaskGenerationApiResponse,
+) {
+  return {
+    title: values.suggested_title ?? "Untitled",
+    description: "",
+    is_saved_task: true,
+    webhook_callback_url: null,
+    proxy_location: "RESIDENTIAL",
+    workflow_definition: {
+      parameters: [
+        {
+          parameter_type: "workflow",
+          workflow_parameter_type: "json",
+          key: "navigation_payload",
+          default_value: JSON.stringify(values.navigation_payload),
+        },
+      ],
+      blocks: [
+        {
+          block_type: "task",
+          label: values.suggested_title,
+          url: values.url,
+          navigation_goal: values.navigation_goal,
+          data_extraction_goal: values.data_extraction_goal,
+          data_schema: values.extracted_information_schema,
+        },
+      ],
+    },
+  };
+}
+
+function createTaskFromTaskGenerationParameters(
+  values: TaskGenerationApiResponse,
+) {
+  return {
+    url: values.url,
+    navigation_goal: values.navigation_goal,
+    data_extraction_goal: values.data_extraction_goal,
+    proxy_location: "RESIDENTIAL",
+    navigation_payload: values.navigation_payload,
+    extracted_information_schema: values.extracted_information_schema,
+  };
+}
+
 function TaskTemplates() {
   const navigate = useNavigate();
   const [prompt, setPrompt] = useState<string>("");
   const credentialGetter = useCredentialGetter();
+  const queryClient = useQueryClient();
 
   const getTaskFromPromptMutation = useMutation({
     mutationFn: async (prompt: string) => {
       const client = await getClient(credentialGetter);
       return client
-        .post("/generate/task", { prompt })
+        .post<
+          { prompt: string },
+          { data: TaskGenerationApiResponse }
+        >("/generate/task", { prompt })
         .then((response) => response.data);
     },
+    onError: (error: AxiosError) => {
+      toast({
+        variant: "destructive",
+        title: "Error creating task from prompt",
+        description: error.message,
+      });
+    },
+  });
+
+  const saveTaskMutation = useMutation({
+    mutationFn: async (params: TaskGenerationApiResponse) => {
+      const client = await getClient(credentialGetter);
+      const templateTask =
+        createTemplateTaskFromTaskGenerationParameters(params);
+      const yaml = convertToYAML(templateTask);
+      return client.post("/workflows", yaml, {
+        headers: {
+          "Content-Type": "text/plain",
+        },
+      });
+    },
+    onSuccess: () => {
+      queryClient.invalidateQueries({
+        queryKey: ["savedTasks"],
+      });
+    },
+    onError: (error: AxiosError) => {
+      toast({
+        variant: "destructive",
+        title: "Error creating task from prompt",
+        description: error.message,
+      });
+    },
+  });
+
+  const runTaskMutation = useMutation({
+    mutationFn: async (params: TaskGenerationApiResponse) => {
+      const client = await getClient(credentialGetter);
+      const data = createTaskFromTaskGenerationParameters(params);
+      return client.post<
+        ReturnType<typeof createTaskFromTaskGenerationParameters>,
+        { data: { task_id: string } }
+      >("/tasks", data);
+    },
     onSuccess: (response) => {
-      navigate("/create/sk-prompt", { state: { data: response } });
+      navigate(`/tasks/${response.data.task_id}/actions`);
     },
     onError: (error: AxiosError) => {
       toast({
@@ -119,13 +214,18 @@ function TaskTemplates() {
             placeholder="Enter your prompt..."
           />
           <div className="h-full">
-            {getTaskFromPromptMutation.isPending ? (
+            {getTaskFromPromptMutation.isPending ||
+            saveTaskMutation.isPending ||
+            runTaskMutation.isPending ? (
               <ReloadIcon className="h-6 w-6 animate-spin" />
             ) : (
               <PaperPlaneIcon
                 className="h-6 w-6 cursor-pointer"
-                onClick={() => {
-                  getTaskFromPromptMutation.mutate(prompt);
+                onClick={async () => {
+                  const taskGenerationResponse =
+                    await getTaskFromPromptMutation.mutateAsync(prompt);
+                  await saveTaskMutation.mutateAsync(taskGenerationResponse);
+                  await runTaskMutation.mutateAsync(taskGenerationResponse);
                 }}
               />
             )}