From 37559fc38f434db3ff15a8fdc7125508ff0c5c7f Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Tue, 15 Oct 2024 08:00:25 -0700
Subject: [PATCH] Render all multi artifacts correctly (#976)

---
 .../src/routes/tasks/detail/HTMLArtifact.tsx  | 53 +++++++++++--------
 .../src/routes/tasks/detail/StepArtifacts.tsx | 20 +++----
 .../src/routes/tasks/detail/TextArtifact.tsx  | 53 +++++++++++--------
 3 files changed, 75 insertions(+), 51 deletions(-)

diff --git a/skyvern-frontend/src/routes/tasks/detail/HTMLArtifact.tsx b/skyvern-frontend/src/routes/tasks/detail/HTMLArtifact.tsx
index ce7d281d..a28a18e4 100644
--- a/skyvern-frontend/src/routes/tasks/detail/HTMLArtifact.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/HTMLArtifact.tsx
@@ -2,7 +2,7 @@ import { artifactApiClient } from "@/api/AxiosClient";
 import { ArtifactApiResponse } from "@/api/types";
 import { Skeleton } from "@/components/ui/skeleton";
 import { Textarea } from "@/components/ui/textarea";
-import { useQuery } from "@tanstack/react-query";
+import { useQueries } from "@tanstack/react-query";
 import axios from "axios";
 
 // https://stackoverflow.com/a/60338028
@@ -27,29 +27,36 @@ function format(html: string) {
 }
 
 type Props = {
-  artifact: ArtifactApiResponse;
+  artifacts: Array<ArtifactApiResponse>;
 };
 
-function HTMLArtifact({ artifact }: Props) {
-  const { data, isFetching, isError, error } = useQuery<string>({
-    queryKey: ["artifact", artifact.artifact_id],
-    queryFn: async () => {
-      if (artifact.uri.startsWith("file://")) {
-        return artifactApiClient
-          .get(`/artifact/text`, {
-            params: {
-              path: artifact.uri.slice(7),
-            },
-          })
-          .then((response) => response.data);
-      }
-      if (artifact.uri.startsWith("s3://") && artifact.signed_url) {
-        return axios.get(artifact.signed_url).then((response) => response.data);
-      }
-    },
+function HTMLArtifact({ artifacts }: Props) {
+  function fetchArtifact(artifact: ArtifactApiResponse) {
+    if (artifact.uri.startsWith("file://")) {
+      return artifactApiClient
+        .get(`/artifact/text`, {
+          params: {
+            path: artifact.uri.slice(7),
+          },
+        })
+        .then((response) => response.data);
+    }
+    if (artifact.uri.startsWith("s3://") && artifact.signed_url) {
+      return axios.get(artifact.signed_url).then((response) => response.data);
+    }
+  }
+
+  const results = useQueries({
+    queries:
+      artifacts?.map((artifact) => {
+        return {
+          queryKey: ["artifact", artifact.artifact_id],
+          queryFn: () => fetchArtifact(artifact),
+        };
+      }) ?? [],
   });
 
-  if (isFetching) {
+  if (results.some((result) => result.isLoading)) {
     return <Skeleton className="h-48 w-full" />;
   }
 
@@ -57,7 +64,11 @@ function HTMLArtifact({ artifact }: Props) {
     <Textarea
       className="w-full"
       rows={15}
-      value={isError ? JSON.stringify(error) : format(data ?? "")}
+      value={
+        results.some((result) => result.isError)
+          ? JSON.stringify(results.find((result) => result.isError)?.error)
+          : results.map((result) => format(result.data ?? "")).join(",\n")
+      }
       readOnly
     />
   );
diff --git a/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx b/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx
index 64a54d49..d263e2e0 100644
--- a/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx
@@ -61,12 +61,12 @@ function StepArtifacts({ id, stepProps }: Props) {
     (artifact) => artifact.artifact_type === ArtifactType.LLMRequest,
   );
 
-  const visibleElementsTreeInPrompt = artifacts?.find(
+  const visibleElementsTreeInPrompt = artifacts?.filter(
     (artifact) =>
       artifact.artifact_type === ArtifactType.VisibleElementsTreeInPrompt,
   );
 
-  const llmPrompt = artifacts?.find(
+  const llmPrompt = artifacts?.filter(
     (artifact) => artifact.artifact_type === ArtifactType.LLMPrompt,
   );
 
@@ -74,7 +74,7 @@ function StepArtifacts({ id, stepProps }: Props) {
     (artifact) => artifact.artifact_type === ArtifactType.LLMResponseParsed,
   );
 
-  const htmlRaw = artifacts?.find(
+  const htmlRaw = artifacts?.filter(
     (artifact) => artifact.artifact_type === ArtifactType.HTMLScrape,
   );
 
@@ -84,8 +84,10 @@ function StepArtifacts({ id, stepProps }: Props) {
         <TabsTrigger value="info">Info</TabsTrigger>
         <TabsTrigger value="screenshot_llm">Annotated Screenshots</TabsTrigger>
         <TabsTrigger value="screenshot_action">Action Screenshots</TabsTrigger>
-        <TabsTrigger value="element_tree_trimmed">Element Tree</TabsTrigger>
-        <TabsTrigger value="html_element_tree">HTML Element Tree</TabsTrigger>
+        <TabsTrigger value="element_tree_trimmed">
+          HTML Element Tree
+        </TabsTrigger>
+        <TabsTrigger value="element_tree">Element Tree</TabsTrigger>
         <TabsTrigger value="llm_prompt">Prompt</TabsTrigger>
         <TabsTrigger value="llm_response_parsed">Action List</TabsTrigger>
         <TabsTrigger value="html_raw">HTML (Raw)</TabsTrigger>
@@ -165,16 +167,16 @@ function StepArtifacts({ id, stepProps }: Props) {
       </TabsContent>
       <TabsContent value="element_tree_trimmed">
         {visibleElementsTreeInPrompt ? (
-          <HTMLArtifact artifact={visibleElementsTreeInPrompt} />
+          <HTMLArtifact artifacts={visibleElementsTreeInPrompt} />
         ) : null}
       </TabsContent>
-      <TabsContent value="html_element_tree">
+      <TabsContent value="element_tree">
         {visibleElementsTree ? (
           <JSONArtifact artifacts={visibleElementsTree} />
         ) : null}
       </TabsContent>
       <TabsContent value="llm_prompt">
-        {llmPrompt ? <TextArtifact artifact={llmPrompt} /> : null}
+        {llmPrompt ? <TextArtifact artifacts={llmPrompt} /> : null}
       </TabsContent>
       <TabsContent value="llm_response_parsed">
         {llmResponseParsed ? (
@@ -182,7 +184,7 @@ function StepArtifacts({ id, stepProps }: Props) {
         ) : null}
       </TabsContent>
       <TabsContent value="html_raw">
-        {htmlRaw ? <TextArtifact artifact={htmlRaw} /> : null}
+        {htmlRaw ? <TextArtifact artifacts={htmlRaw} /> : null}
       </TabsContent>
       <TabsContent value="llm_request">
         {llmRequest ? <JSONArtifact artifacts={llmRequest} /> : null}
diff --git a/skyvern-frontend/src/routes/tasks/detail/TextArtifact.tsx b/skyvern-frontend/src/routes/tasks/detail/TextArtifact.tsx
index 9e626df9..e154a87c 100644
--- a/skyvern-frontend/src/routes/tasks/detail/TextArtifact.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/TextArtifact.tsx
@@ -2,33 +2,40 @@ import { artifactApiClient } from "@/api/AxiosClient";
 import { ArtifactApiResponse } from "@/api/types";
 import { Skeleton } from "@/components/ui/skeleton";
 import { Textarea } from "@/components/ui/textarea";
-import { useQuery } from "@tanstack/react-query";
+import { useQueries } from "@tanstack/react-query";
 import axios from "axios";
 
 type Props = {
-  artifact: ArtifactApiResponse;
+  artifacts: Array<ArtifactApiResponse>;
 };
 
-function TextArtifact({ artifact }: Props) {
-  const { data, isFetching, isError, error } = useQuery<string>({
-    queryKey: ["artifact", artifact.artifact_id],
-    queryFn: async () => {
-      if (artifact.uri.startsWith("file://")) {
-        return artifactApiClient
-          .get(`/artifact/text`, {
-            params: {
-              path: artifact.uri.slice(7),
-            },
-          })
-          .then((response) => response.data);
-      }
-      if (artifact.uri.startsWith("s3://") && artifact.signed_url) {
-        return axios.get(artifact.signed_url).then((response) => response.data);
-      }
-    },
+function TextArtifact({ artifacts }: Props) {
+  function fetchArtifact(artifact: ArtifactApiResponse) {
+    if (artifact.uri.startsWith("file://")) {
+      return artifactApiClient
+        .get(`/artifact/text`, {
+          params: {
+            path: artifact.uri.slice(7),
+          },
+        })
+        .then((response) => response.data);
+    }
+    if (artifact.uri.startsWith("s3://") && artifact.signed_url) {
+      return axios.get(artifact.signed_url).then((response) => response.data);
+    }
+  }
+
+  const results = useQueries({
+    queries:
+      artifacts?.map((artifact) => {
+        return {
+          queryKey: ["artifact", artifact.artifact_id],
+          queryFn: () => fetchArtifact(artifact),
+        };
+      }) ?? [],
   });
 
-  if (isFetching) {
+  if (results.some((result) => result.isLoading)) {
     return <Skeleton className="h-48 w-full" />;
   }
 
@@ -36,7 +43,11 @@ function TextArtifact({ artifact }: Props) {
     <Textarea
       className="w-full"
       rows={15}
-      value={isError ? JSON.stringify(error) : data}
+      value={
+        results.some((result) => result.isError)
+          ? JSON.stringify(results.find((result) => result.isError)?.error)
+          : results.map((result) => result.data).join(",\n")
+      }
       readOnly
     />
   );