From 4e436395315c4259dfeba4880c8c106a94b8b726 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Fri, 27 Dec 2024 18:15:40 -0800
Subject: [PATCH] add workflow run block screenshots (#1443)

---
 .../workflowRun/ObserverThoughtScreenshot.tsx |  2 +-
 .../WorkflowRunBlockScreenshot.tsx            | 66 +++++++++++++++++++
 .../workflowRun/WorkflowRunOverview.tsx       |  7 +-
 3 files changed, 71 insertions(+), 4 deletions(-)
 create mode 100644 skyvern-frontend/src/routes/workflows/workflowRun/WorkflowRunBlockScreenshot.tsx

diff --git a/skyvern-frontend/src/routes/workflows/workflowRun/ObserverThoughtScreenshot.tsx b/skyvern-frontend/src/routes/workflows/workflowRun/ObserverThoughtScreenshot.tsx
index 705e7321..5daa2389 100644
--- a/skyvern-frontend/src/routes/workflows/workflowRun/ObserverThoughtScreenshot.tsx
+++ b/skyvern-frontend/src/routes/workflows/workflowRun/ObserverThoughtScreenshot.tsx
@@ -61,7 +61,7 @@ function ObserverThoughtScreenshot({ observerThoughtId, taskStatus }: Props) {
   if (!screenshot) {
     return (
       <div className="flex h-full items-center justify-center bg-slate-elevation1">
-        No screenshot found for this action.
+        No screenshot found for this thought.
       </div>
     );
   }
diff --git a/skyvern-frontend/src/routes/workflows/workflowRun/WorkflowRunBlockScreenshot.tsx b/skyvern-frontend/src/routes/workflows/workflowRun/WorkflowRunBlockScreenshot.tsx
new file mode 100644
index 00000000..4ac6bb6b
--- /dev/null
+++ b/skyvern-frontend/src/routes/workflows/workflowRun/WorkflowRunBlockScreenshot.tsx
@@ -0,0 +1,66 @@
+import { getClient } from "@/api/AxiosClient";
+import { ArtifactApiResponse, ArtifactType } from "@/api/types";
+import { ZoomableImage } from "@/components/ZoomableImage";
+import { useCredentialGetter } from "@/hooks/useCredentialGetter";
+import { useQuery } from "@tanstack/react-query";
+import { ReloadIcon } from "@radix-ui/react-icons";
+import { getImageURL } from "@/routes/tasks/detail/artifactUtils";
+
+type Props = {
+  workflowRunBlockId: string;
+};
+
+function WorkflowRunBlockScreenshot({ workflowRunBlockId }: Props) {
+  const credentialGetter = useCredentialGetter();
+
+  const { data: artifacts, isLoading } = useQuery<Array<ArtifactApiResponse>>({
+    queryKey: ["workflowRunBlock", workflowRunBlockId, "artifacts"],
+    queryFn: async () => {
+      const client = await getClient(credentialGetter);
+      return client
+        .get(`/workflow_run_block/${workflowRunBlockId}/artifacts`)
+        .then((response) => response.data);
+    },
+    refetchInterval: (query) => {
+      const data = query.state.data;
+      const screenshot = data?.filter(
+        (artifact) => artifact.artifact_type === ArtifactType.LLMScreenshot,
+      )?.[0];
+      if (!screenshot) {
+        return 5000;
+      }
+      return false;
+    },
+  });
+
+  const llmScreenshots = artifacts?.filter(
+    (artifact) => artifact.artifact_type === ArtifactType.LLMScreenshot,
+  );
+
+  const screenshot = llmScreenshots?.[0];
+
+  if (isLoading) {
+    return (
+      <div className="flex h-full items-center justify-center gap-2 bg-slate-elevation1">
+        <ReloadIcon className="h-6 w-6 animate-spin" />
+        <div>Loading screenshot...</div>
+      </div>
+    );
+  }
+
+  if (!screenshot) {
+    return (
+      <div className="flex h-full items-center justify-center bg-slate-elevation1">
+        No screenshot found for this workflow run block.
+      </div>
+    );
+  }
+
+  return (
+    <figure className="mx-auto flex max-w-full flex-col items-center gap-2 overflow-hidden rounded">
+      <ZoomableImage src={getImageURL(screenshot)} alt="llm-screenshot" />
+    </figure>
+  );
+}
+
+export { WorkflowRunBlockScreenshot };
diff --git a/skyvern-frontend/src/routes/workflows/workflowRun/WorkflowRunOverview.tsx b/skyvern-frontend/src/routes/workflows/workflowRun/WorkflowRunOverview.tsx
index bfa655b2..143c9838 100644
--- a/skyvern-frontend/src/routes/workflows/workflowRun/WorkflowRunOverview.tsx
+++ b/skyvern-frontend/src/routes/workflows/workflowRun/WorkflowRunOverview.tsx
@@ -23,6 +23,7 @@ import { DotFilledIcon } from "@radix-ui/react-icons";
 import { WorkflowRunTimelineItemInfoSection } from "./WorkflowRunTimelineItemInfoSection";
 import { ObserverThoughtScreenshot } from "./ObserverThoughtScreenshot";
 import { ScrollArea, ScrollAreaViewport } from "@/components/ui/scroll-area";
+import { WorkflowRunBlockScreenshot } from "./WorkflowRunBlockScreenshot";
 
 export type ActionItem = {
   block: WorkflowRunBlock;
@@ -103,9 +104,9 @@ function WorkflowRunOverview() {
             />
           )}
           {isWorkflowRunBlock(selection) && (
-            <div className="flex h-full w-full items-center justify-center bg-slate-elevation1">
-              No screenshot found for this block
-            </div>
+            <WorkflowRunBlockScreenshot
+              workflowRunBlockId={selection.workflow_run_block_id}
+            />
           )}
           {isObserverThought(selection) && (
             <ObserverThoughtScreenshot