From 240e5912808482644eddff65ad32bab9ff1505c6 Mon Sep 17 00:00:00 2001
From: Kerem Yilmaz <kerem@skyvern.com>
Date: Thu, 11 Jul 2024 08:03:18 -0700
Subject: [PATCH] Show steps in task page (#585)

---
 poetry.lock                                   | 18 ++--
 pyproject.toml                                |  2 +-
 .../src/hooks/useCostCalculator.ts            |  9 ++
 .../tasks/detail/StepArtifactsLayout.tsx      |  6 +-
 .../routes/tasks/detail/StepNavigation.tsx    |  6 +-
 .../src/routes/tasks/detail/TaskDetails.tsx   | 36 ++------
 .../src/routes/tasks/detail/TaskInfo.tsx      | 90 +++++++++++++++++++
 .../routes/tasks/detail/hooks/useTaskQuery.ts | 33 +++++++
 .../src/store/CostCalculatorContext.ts        |  8 ++
 9 files changed, 157 insertions(+), 51 deletions(-)
 create mode 100644 skyvern-frontend/src/hooks/useCostCalculator.ts
 create mode 100644 skyvern-frontend/src/routes/tasks/detail/TaskInfo.tsx
 create mode 100644 skyvern-frontend/src/routes/tasks/detail/hooks/useTaskQuery.ts
 create mode 100644 skyvern-frontend/src/store/CostCalculatorContext.ts

diff --git a/poetry.lock b/poetry.lock
index 9e714594..e79be490 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -2943,13 +2943,13 @@ files = [
 
 [[package]]
 name = "litellm"
-version = "1.41.12"
+version = "1.35.35"
 description = "Library to easily interface with LLM API providers"
 optional = false
 python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8"
 files = [
-    {file = "litellm-1.41.12-py3-none-any.whl", hash = "sha256:9af8b65ca48f0aa5b8ef10a63c21b00553843aa1e498f2c9308738b97d4a50f3"},
-    {file = "litellm-1.41.12.tar.gz", hash = "sha256:f94b5ac8857ea8b98b87f7d3071dbd10e24fe9e0d7831969adafb549688ce0ce"},
+    {file = "litellm-1.35.35-py3-none-any.whl", hash = "sha256:9dcf4a1fbdb1864de03f2fbcc26cc19ee83449966ad0b373b1d61f64ade963f7"},
+    {file = "litellm-1.35.35.tar.gz", hash = "sha256:f8873892f4a2f082e2f5f4fed5740f341b7d1a7778445785b2af68adbc2793e9"},
 ]
 
 [package.dependencies]
@@ -2957,17 +2957,15 @@ aiohttp = "*"
 click = "*"
 importlib-metadata = ">=6.8.0"
 jinja2 = ">=3.1.2,<4.0.0"
-jsonschema = ">=4.22.0,<5.0.0"
-openai = ">=1.27.0"
-pydantic = ">=2.0.0,<3.0.0"
+openai = ">=1.0.0"
 python-dotenv = ">=0.2.0"
 requests = ">=2.31.0,<3.0.0"
-tiktoken = ">=0.7.0"
+tiktoken = ">=0.4.0"
 tokenizers = "*"
 
 [package.extras]
-extra-proxy = ["azure-identity (>=1.15.0,<2.0.0)", "azure-keyvault-secrets (>=4.8.0,<5.0.0)", "google-cloud-kms (>=2.21.3,<3.0.0)", "prisma (==0.11.0)", "pynacl (>=1.5.0,<2.0.0)", "resend (>=0.8.0,<0.9.0)"]
-proxy = ["PyJWT (>=2.8.0,<3.0.0)", "apscheduler (>=3.10.4,<4.0.0)", "backoff", "cryptography (>=42.0.5,<43.0.0)", "fastapi (>=0.111.0,<0.112.0)", "fastapi-sso (>=0.10.0,<0.11.0)", "gunicorn (>=22.0.0,<23.0.0)", "orjson (>=3.9.7,<4.0.0)", "python-multipart (>=0.0.9,<0.0.10)", "pyyaml (>=6.0.1,<7.0.0)", "rq", "uvicorn (>=0.22.0,<0.23.0)"]
+extra-proxy = ["azure-identity (>=1.15.0,<2.0.0)", "azure-keyvault-secrets (>=4.8.0,<5.0.0)", "google-cloud-kms (>=2.21.3,<3.0.0)", "prisma (==0.11.0)", "resend (>=0.8.0,<0.9.0)"]
+proxy = ["PyJWT (>=2.8.0,<3.0.0)", "apscheduler (>=3.10.4,<4.0.0)", "backoff", "cryptography (>=42.0.5,<43.0.0)", "fastapi (>=0.109.1,<0.110.0)", "fastapi-sso (>=0.10.0,<0.11.0)", "gunicorn (>=21.2.0,<22.0.0)", "orjson (>=3.9.7,<4.0.0)", "python-multipart (>=0.0.9,<0.0.10)", "pyyaml (>=6.0.1,<7.0.0)", "rq", "uvicorn (>=0.22.0,<0.23.0)"]
 
 [[package]]
 name = "lxml"
@@ -7277,4 +7275,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11,<3.12"
-content-hash = "40e1dba8128bed75196f0c8dc21894f1b37f76a16f19b1678697df7563f67a7a"
+content-hash = "b3bb150f1811469740857a9bf062345a9597ace0d292b83f08a9f3e0707519d8"
diff --git a/pyproject.toml b/pyproject.toml
index 835e5ebf..9c5ec0b9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,7 +19,7 @@ python-multipart = "^0.0.6"
 toml = "^0.10.2"
 jinja2 = "^3.1.2"
 uvicorn = {extras = ["standard"], version = "^0.24.0.post1"}
-litellm = "1.41.12"
+litellm = "1.35.35"
 duckduckgo-search = "^3.8.0"
 selenium = "^4.13.0"
 bs4 = "^0.0.1"
diff --git a/skyvern-frontend/src/hooks/useCostCalculator.ts b/skyvern-frontend/src/hooks/useCostCalculator.ts
new file mode 100644
index 00000000..a0582507
--- /dev/null
+++ b/skyvern-frontend/src/hooks/useCostCalculator.ts
@@ -0,0 +1,9 @@
+import { CostCalculatorContext } from "@/store/CostCalculatorContext";
+import { useContext } from "react";
+
+function useCostCalculator() {
+  const costCalculator = useContext(CostCalculatorContext);
+  return costCalculator;
+}
+
+export { useCostCalculator };
diff --git a/skyvern-frontend/src/routes/tasks/detail/StepArtifactsLayout.tsx b/skyvern-frontend/src/routes/tasks/detail/StepArtifactsLayout.tsx
index a024a8a2..e9addf75 100644
--- a/skyvern-frontend/src/routes/tasks/detail/StepArtifactsLayout.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/StepArtifactsLayout.tsx
@@ -30,11 +30,7 @@ function StepArtifactsLayout() {
     return <div>Error: {error?.message}</div>;
   }
 
-  if (!steps) {
-    return <div>No steps found</div>;
-  }
-
-  const activeStep = steps[activeIndex];
+  const activeStep = steps?.[activeIndex];
 
   return (
     <div className="flex">
diff --git a/skyvern-frontend/src/routes/tasks/detail/StepNavigation.tsx b/skyvern-frontend/src/routes/tasks/detail/StepNavigation.tsx
index 77b57955..ef6f9f8d 100644
--- a/skyvern-frontend/src/routes/tasks/detail/StepNavigation.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/StepNavigation.tsx
@@ -41,13 +41,9 @@ function StepNavigation({ activeIndex, onActiveIndexChange }: Props) {
     return <div>Error: {error?.message}</div>;
   }
 
-  if (!steps) {
-    return <div>No steps found</div>;
-  }
-
   return (
     <nav className="flex flex-col gap-4">
-      {steps.map((step, index) => {
+      {steps?.map((step, index) => {
         const isActive = activeIndex === index;
         return (
           <div
diff --git a/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx b/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx
index 68d215f2..31486908 100644
--- a/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx
@@ -1,6 +1,5 @@
 import { getClient } from "@/api/AxiosClient";
-import { Status, TaskApiResponse } from "@/api/types";
-import { StatusBadge } from "@/components/StatusBadge";
+import { Status } from "@/api/types";
 import { Button } from "@/components/ui/button";
 import {
   Dialog,
@@ -19,13 +18,10 @@ import { toast } from "@/components/ui/use-toast";
 import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 import { cn } from "@/util/utils";
 import { ReloadIcon } from "@radix-ui/react-icons";
-import {
-  keepPreviousData,
-  useMutation,
-  useQuery,
-  useQueryClient,
-} from "@tanstack/react-query";
+import { useMutation, useQueryClient } from "@tanstack/react-query";
 import { NavLink, Outlet, useParams } from "react-router-dom";
+import { TaskInfo } from "./TaskInfo";
+import { useTaskQuery } from "./hooks/useTaskQuery";
 
 function TaskDetails() {
   const { taskId } = useParams();
@@ -37,23 +33,7 @@ function TaskDetails() {
     isLoading: taskIsLoading,
     isError: taskIsError,
     error: taskError,
-  } = useQuery<TaskApiResponse>({
-    queryKey: ["task", taskId],
-    queryFn: async () => {
-      const client = await getClient(credentialGetter);
-      return client.get(`/tasks/${taskId}`).then((response) => response.data);
-    },
-    refetchInterval: (query) => {
-      if (
-        query.state.data?.status === Status.Running ||
-        query.state.data?.status === Status.Queued
-      ) {
-        return 10000;
-      }
-      return false;
-    },
-    placeholderData: keepPreviousData,
-  });
+  } = useTaskQuery({ id: taskId });
 
   const cancelTaskMutation = useMutation({
     mutationFn: async () => {
@@ -124,11 +104,7 @@ function TaskDetails() {
       <div className="flex justify-between items-center">
         <div className="flex items-center gap-4">
           <span className="text-lg">{taskId}</span>
-          {taskIsLoading ? (
-            <Skeleton className="w-28 h-8" />
-          ) : task ? (
-            <StatusBadge status={task?.status} />
-          ) : null}
+          {taskId && <TaskInfo id={taskId} />}
         </div>
         {taskIsRunningOrQueued && (
           <Dialog>
diff --git a/skyvern-frontend/src/routes/tasks/detail/TaskInfo.tsx b/skyvern-frontend/src/routes/tasks/detail/TaskInfo.tsx
new file mode 100644
index 00000000..7ff000df
--- /dev/null
+++ b/skyvern-frontend/src/routes/tasks/detail/TaskInfo.tsx
@@ -0,0 +1,90 @@
+import { Badge } from "@/components/ui/badge";
+import { StatusBadge } from "@/components/StatusBadge";
+import { Skeleton } from "@/components/ui/skeleton";
+import { Status, StepApiResponse } from "@/api/types";
+import { keepPreviousData, useQuery } from "@tanstack/react-query";
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
+import { getClient } from "@/api/AxiosClient";
+import { useTaskQuery } from "./hooks/useTaskQuery";
+import { useCostCalculator } from "@/hooks/useCostCalculator";
+
+type Props = {
+  id: string;
+};
+
+const formatter = Intl.NumberFormat("en-US", {
+  style: "currency",
+  currency: "USD",
+});
+
+function TaskInfo({ id }: Props) {
+  const credentialGetter = useCredentialGetter();
+  const costCalculator = useCostCalculator();
+  const {
+    data: task,
+    isLoading: taskIsLoading,
+    isError: taskIsError,
+  } = useTaskQuery({ id });
+
+  const taskIsRunningOrQueued =
+    task?.status === Status.Running || task?.status === Status.Queued;
+
+  const {
+    data: steps,
+    isLoading: stepsIsLoading,
+    isError: stepsIsError,
+  } = useQuery<Array<StepApiResponse>>({
+    queryKey: ["task", id, "steps"],
+    queryFn: async () => {
+      const client = await getClient(credentialGetter);
+      return client.get(`/tasks/${id}/steps`).then((response) => response.data);
+    },
+    refetchOnWindowFocus: taskIsRunningOrQueued,
+    refetchInterval: taskIsRunningOrQueued ? 5000 : false,
+    placeholderData: keepPreviousData,
+  });
+
+  if (stepsIsLoading || taskIsLoading) {
+    return (
+      <div className="flex gap-4">
+        <Skeleton className="w-20 h-6" />
+        <Skeleton className="w-20 h-6" />
+        <Skeleton className="w-20 h-6" />
+        <Skeleton className="w-20 h-6" />
+      </div>
+    );
+  }
+
+  if (stepsIsError || taskIsError) {
+    return null;
+  }
+
+  const actionCount = steps?.reduce((acc, step) => {
+    const actionsAndResults = step.output?.actions_and_results ?? [];
+
+    const actionCount = actionsAndResults.reduce((acc, actionAndResult) => {
+      const actionResult = actionAndResult[1];
+      if (actionResult.length === 0) {
+        return acc;
+      }
+      return acc + 1;
+    }, 0);
+    return acc + actionCount;
+  }, 0);
+
+  const showCost = typeof costCalculator === "function";
+  const notRunningSteps = steps?.filter((step) => step.status !== "running");
+
+  return (
+    <div className="flex gap-4">
+      {task && <StatusBadge status={task.status} />}
+      <Badge>Steps: {notRunningSteps?.length}</Badge>
+      <Badge>Actions: {actionCount}</Badge>
+      {showCost && (
+        <Badge>Cost: {formatter.format(costCalculator(steps ?? []))}</Badge>
+      )}
+    </div>
+  );
+}
+
+export { TaskInfo };
diff --git a/skyvern-frontend/src/routes/tasks/detail/hooks/useTaskQuery.ts b/skyvern-frontend/src/routes/tasks/detail/hooks/useTaskQuery.ts
new file mode 100644
index 00000000..b4e19a57
--- /dev/null
+++ b/skyvern-frontend/src/routes/tasks/detail/hooks/useTaskQuery.ts
@@ -0,0 +1,33 @@
+import { getClient } from "@/api/AxiosClient";
+import { Status, TaskApiResponse } from "@/api/types";
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
+import { keepPreviousData, useQuery } from "@tanstack/react-query";
+
+type Props = {
+  id?: string;
+};
+
+function useTaskQuery({ id }: Props) {
+  const credentialGetter = useCredentialGetter();
+
+  return useQuery<TaskApiResponse>({
+    queryKey: ["task", id],
+    queryFn: async () => {
+      const client = await getClient(credentialGetter);
+      return client.get(`/tasks/${id}`).then((response) => response.data);
+    },
+    enabled: !!id,
+    refetchInterval: (query) => {
+      if (
+        query.state.data?.status === Status.Running ||
+        query.state.data?.status === Status.Queued
+      ) {
+        return 5000;
+      }
+      return false;
+    },
+    placeholderData: keepPreviousData,
+  });
+}
+
+export { useTaskQuery };
diff --git a/skyvern-frontend/src/store/CostCalculatorContext.ts b/skyvern-frontend/src/store/CostCalculatorContext.ts
new file mode 100644
index 00000000..e72c9a13
--- /dev/null
+++ b/skyvern-frontend/src/store/CostCalculatorContext.ts
@@ -0,0 +1,8 @@
+import { StepApiResponse } from "@/api/types";
+import { createContext } from "react";
+
+type TaskCostCalculator = (steps: Array<StepApiResponse>) => number;
+
+const CostCalculatorContext = createContext<TaskCostCalculator | null>(null);
+
+export { CostCalculatorContext };