From 4ab90aa3fa02ac3e6884fb3cf035e80841ca4c22 Mon Sep 17 00:00:00 2001
From: Flavien David <flavien.david74@gmail.com>
Date: Thu, 25 Jul 2024 12:14:19 +0200
Subject: [PATCH] Flav/viz draft (#6486)

* Working iframe

Saving progress, dirty code,

File transfer working.

Working. Need UI work + code clean up.

Build working

* Fix file upload on front edge. to be reverted.

* Fix UX

* Tab manually changed.

* Fix UI and rendering.

* Better naming

* Clean up

* Adding front-viz

* Adding comments.

* Clean up console.log

* Remove unused image.

* Draft

* Create .nvmrc

* Add foundations for viz

* :sparkles:

* :see_no_evil:

* :see_no_evil:

* Fix

* Fix useFile hook

* :scissors:

* :scissors:

* Remove Sparkle from viz

* :sparkles:

* :back:

* Use NEXT_PUBLIC_VIZ_URL at build time

* :see_no_evil:

* :scissors:

* :see_no_evil:

* :scissors:

* :expressionless:

* :monkey_face:

* :hammer:

* :boom:

* Only listen to current action

* :scissors:

* :scissors:

* :sparkles:

* :scissors:

---------

Co-authored-by: Aric Lasry <lasry.aric@gmail.com>
---
 front/Dockerfile                              |   5 +-
 .../assistant/conversation/AgentMessage.tsx   |  56 +++---
 .../actions/VisualizationActionIframe.tsx     | 171 ++++++++++++++++++
 .../submitAssistantBuilderForm.ts             |   2 +-
 .../api/assistant/actions/visualization.ts    | 111 ++++++++++--
 front/lib/api/assistant/generation.ts         |   9 +-
 .../resources/content_fragment_resource.ts    |   6 +-
 front/pages/api/w/[wId]/files/[fileId].ts     |   2 +-
 k8s/cloudbuild.yaml                           |   9 +-
 .../front/assistant/actions/visualization.ts  | 125 ++++++++++++-
 types/src/front/lib/actions/registry.ts       |   2 +-
 viz/app/components/VisualizationWrapper.tsx   |  41 +++--
 12 files changed, 457 insertions(+), 82 deletions(-)
 create mode 100644 front/components/assistant/conversation/actions/VisualizationActionIframe.tsx

diff --git a/front/Dockerfile b/front/Dockerfile
index 9934dceb7709..fc472b9fd3b9 100644
--- a/front/Dockerfile
+++ b/front/Dockerfile
@@ -16,7 +16,10 @@ RUN npm ci
 COPY /front .
 
 ARG COMMIT_HASH
-ENV NEXT_PUBLIC_COMMIT_HASH=${COMMIT_HASH}
+ARG NEXT_PUBLIC_VIZ_URL
+
+ENV NEXT_PUBLIC_COMMIT_HASH=$COMMIT_HASH
+ENV NEXT_PUBLIC_VIZ_URL=$NEXT_PUBLIC_VIZ_URL
 
 # fake database URIs are needed because Sequelize will throw if the `url` parameter
 # is undefined, and `next build` imports the `models.ts` file while "Collecting page data"
diff --git a/front/components/assistant/conversation/AgentMessage.tsx b/front/components/assistant/conversation/AgentMessage.tsx
index 244f3ce4c039..66048b422646 100644
--- a/front/components/assistant/conversation/AgentMessage.tsx
+++ b/front/components/assistant/conversation/AgentMessage.tsx
@@ -1,6 +1,5 @@
 import {
   ArrowPathIcon,
-  BracesIcon,
   Button,
   Chip,
   Citation,
@@ -22,7 +21,6 @@ import type {
   LightAgentConfigurationType,
   RetrievalActionType,
   UserType,
-  VisualizationActionType,
   WebsearchActionType,
   WebsearchResultType,
   WorkspaceType,
@@ -39,6 +37,7 @@ import {
   isWebsearchActionType,
   removeNulls,
 } from "@dust-tt/types";
+import assert from "assert";
 import Link from "next/link";
 import { useRouter } from "next/router";
 import { useCallback, useContext, useEffect, useRef, useState } from "react";
@@ -46,6 +45,7 @@ import { useCallback, useContext, useEffect, useRef, useState } from "react";
 import { makeDocumentCitations } from "@app/components/actions/retrieval/utils";
 import { AssistantDetailsDropdownMenu } from "@app/components/assistant/AssistantDetailsDropdownMenu";
 import { AgentMessageActions } from "@app/components/assistant/conversation/actions/AgentMessageActions";
+import { VisualizationActionIframe } from "@app/components/assistant/conversation/actions/VisualizationActionIframe";
 import type { MessageSizeType } from "@app/components/assistant/conversation/ConversationMessage";
 import { ConversationMessage } from "@app/components/assistant/conversation/ConversationMessage";
 import { GenerationContext } from "@app/components/assistant/conversation/GenerationContextProvider";
@@ -88,19 +88,9 @@ export function AgentMessage({
   const [streamedAgentMessage, setStreamedAgentMessage] =
     useState<AgentMessageType>(message);
 
-  const defaultVisualizations: VisualizationActionType[] =
-    message.actions.filter((a): a is VisualizationActionType =>
-      isVisualizationActionType(a)
-    ) as VisualizationActionType[];
-
   const [streamedVisualizations, setStreamedVisualizations] = useState<
     { actionId: number; visualization: string }[]
-  >(
-    defaultVisualizations.map((v) => ({
-      actionId: v.id,
-      visualization: v.generation ?? "",
-    }))
-  );
+  >([]);
 
   const [isRetryHandlerProcessing, setIsRetryHandlerProcessing] =
     useState<boolean>(false);
@@ -213,6 +203,7 @@ export function AgentMessage({
             ...event.message,
           };
         });
+        setStreamedVisualizations([]);
         break;
       }
 
@@ -535,24 +526,6 @@ export function AgentMessage({
           </div>
         ) : null}
 
-        {/* This is where we will we plug Aric's work to render the graph in an iframe. */}
-        {streamedVisualizations.map(({ actionId, visualization }) => {
-          return (
-            <div key={actionId}>
-              <div className="flex flex-row gap-2">
-                <Icon size="sm" visual={BracesIcon} />
-                <div className="font-semibold">Visualization</div>
-              </div>
-              <div>
-                <RenderMessageMarkdown
-                  content={"```js" + visualization + "\n```"}
-                  isStreaming={true}
-                />
-              </div>
-            </div>
-          );
-        })}
-
         {agentMessage.content !== null && (
           <div>
             {lastTokenClassification !== "chain_of_thought" &&
@@ -583,6 +556,27 @@ export function AgentMessage({
             )}
           </div>
         )}
+        <>
+          {agentMessage.actions
+            .filter((a) => isVisualizationActionType(a))
+            .map((a, i) => {
+              const streamingViz = streamedVisualizations.find(
+                (sv) => sv.actionId === a.id
+              );
+              assert(isVisualizationActionType(a));
+              return (
+                <VisualizationActionIframe
+                  action={a}
+                  conversationId={conversationId}
+                  isStreaming={!!streamingViz}
+                  key={i}
+                  onRetry={() => retryHandler(agentMessage)}
+                  owner={owner}
+                  streamedCode={streamingViz?.visualization || null}
+                />
+              );
+            })}
+        </>
         {agentMessage.status === "cancelled" && (
           <Chip
             label="Message generation was interrupted"
diff --git a/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx b/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
new file mode 100644
index 000000000000..175ee7f2b4cf
--- /dev/null
+++ b/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
@@ -0,0 +1,171 @@
+import { BracesIcon, PlayIcon, Tab } from "@dust-tt/sparkle";
+import type {
+  VisualizationActionType,
+  VisualizationRPCRequest,
+  WorkspaceType,
+} from "@dust-tt/types";
+import {
+  isGetCodeToExecuteRequest,
+  isGetFileRequest,
+  isVisualizationRPCRequest,
+  visualizationExtractCodeNonStreaming,
+  visualizationExtractCodeStreaming,
+} from "@dust-tt/types";
+import { useCallback, useEffect, useState } from "react";
+
+import { RenderMessageMarkdown } from "@app/components/assistant/RenderMessageMarkdown";
+
+const sendResponseToIframe = (
+  request: VisualizationRPCRequest,
+  response: unknown,
+  target: MessageEventSource
+) => {
+  target.postMessage(
+    {
+      command: "answer",
+      messageUniqueId: request.messageUniqueId,
+      actionId: request.actionId,
+      result: response,
+    },
+    // TODO(2024-07-24 flav) Restrict origin.
+    { targetOrigin: "*" }
+  );
+};
+
+// Custom hook to encapsulate the logic for handling visualization messages.
+function useVisualizationDataHandler(
+  action: VisualizationActionType,
+  workspaceId: string,
+  onRetry: () => void
+) {
+  const getFile = useCallback(
+    async (fileId: string) => {
+      const response = await fetch(
+        `/api/w/${workspaceId}/files/${fileId}?action=view`
+      );
+      if (!response.ok) {
+        // TODO(2024-07-24 flav) Propagate the error to the iframe.
+        throw new Error(`Failed to fetch file ${fileId}`);
+      }
+
+      const resBuffer = await response.arrayBuffer();
+      return new File([resBuffer], fileId, {
+        type: response.headers.get("Content-Type") || undefined,
+      });
+    },
+    [workspaceId]
+  );
+
+  useEffect(() => {
+    const listener = async (event: MessageEvent) => {
+      const { data } = event;
+
+      // TODO(2024-07-24 flav) Check origin.
+      if (
+        !isVisualizationRPCRequest(data) ||
+        !event.source ||
+        data.actionId !== action.id
+      ) {
+        return;
+      }
+
+      if (isGetFileRequest(data)) {
+        const file = await getFile(data.params.fileId);
+
+        sendResponseToIframe(data, { file }, event.source);
+      } else if (isGetCodeToExecuteRequest(data)) {
+        const code = action.generation;
+
+        sendResponseToIframe(data, { code }, event.source);
+      } else {
+        // TODO(2024-07-24 flav) Pass the error message to the host window.
+        onRetry();
+      }
+
+      // TODO: Types above are not accurate, as it can pass the first check but won't enter any if block.
+    };
+
+    window.addEventListener("message", listener);
+    return () => window.removeEventListener("message", listener);
+  }, [action.generation, action.id, onRetry, getFile]);
+
+  return { getFile };
+}
+
+export function VisualizationActionIframe({
+  owner,
+  action,
+  isStreaming,
+  streamedCode,
+  onRetry,
+}: {
+  conversationId: string;
+  owner: WorkspaceType;
+  action: VisualizationActionType;
+  streamedCode: string | null;
+  isStreaming: boolean;
+  onRetry: () => void;
+}) {
+  const [activeTab, setActiveTab] = useState<"code" | "runtime">("code");
+  const [tabManuallyChanged, setTabManuallyChanged] = useState(false);
+
+  const workspaceId = owner.sId;
+
+  useVisualizationDataHandler(action, workspaceId, onRetry);
+
+  useEffect(() => {
+    if (activeTab === "code" && action.generation && !tabManuallyChanged) {
+      setActiveTab("runtime");
+      setTabManuallyChanged(true);
+    }
+  }, [action.generation, activeTab, tabManuallyChanged]);
+
+  let extractedCode: string | null = null;
+
+  if (action.generation) {
+    extractedCode = visualizationExtractCodeNonStreaming(action.generation);
+  } else {
+    extractedCode = visualizationExtractCodeStreaming(streamedCode || "");
+  }
+
+  return (
+    <>
+      <Tab
+        tabs={[
+          {
+            label: "Code",
+            id: "code",
+            current: activeTab === "code",
+            icon: BracesIcon,
+            sizing: "expand",
+          },
+          {
+            label: "Run",
+            id: "runtime",
+            current: activeTab === "runtime",
+            icon: PlayIcon,
+            sizing: "expand",
+            hasSeparator: true,
+          },
+        ]}
+        setCurrentTab={(tabId, event) => {
+          event.preventDefault();
+          setActiveTab(tabId as "code" | "runtime");
+        }}
+      />
+      {activeTab === "code" && extractedCode && extractedCode.length > 0 && (
+        <RenderMessageMarkdown
+          content={"```javascript\n" + extractedCode + "\n```"}
+          isStreaming={isStreaming}
+        />
+      )}
+      {activeTab === "runtime" && (
+        <iframe
+          style={{ width: "100%", height: "600px" }}
+          src={`${process.env.NEXT_PUBLIC_VIZ_URL}/content?aId=${action.id}`}
+          sandbox="allow-scripts"
+        />
+      )}
+    </>
+  );
+}
diff --git a/front/components/assistant_builder/submitAssistantBuilderForm.ts b/front/components/assistant_builder/submitAssistantBuilderForm.ts
index cebeebd6593b..8bf6de3dff85 100644
--- a/front/components/assistant_builder/submitAssistantBuilderForm.ts
+++ b/front/components/assistant_builder/submitAssistantBuilderForm.ts
@@ -178,7 +178,7 @@ export async function submitAssistantBuilderForm({
           {
             type: "visualization_configuration",
             name: DEFAULT_VISUALIZATION_ACTION_NAME,
-            description: "Generates graphs from data.",
+            description: "Generate client side javascript react code.",
           },
         ];
 
diff --git a/front/lib/api/assistant/actions/visualization.ts b/front/lib/api/assistant/actions/visualization.ts
index 30e313652ea0..e65e82b8dc00 100644
--- a/front/lib/api/assistant/actions/visualization.ts
+++ b/front/lib/api/assistant/actions/visualization.ts
@@ -17,11 +17,15 @@ import {
   CLAUDE_3_5_SONNET_DEFAULT_MODEL_CONFIG,
   cloneBaseConfig,
   DustProdActionRegistry,
+  isContentFragmentType,
   isProviderWhitelisted,
   Ok,
   VisualizationActionOutputSchema,
 } from "@dust-tt/types";
+import assert from "assert";
 import { isLeft } from "fp-ts/lib/Either";
+import readline from "readline";
+import type { Readable } from "stream";
 
 import { runActionStreamed } from "@app/lib/actions/server";
 import { DEFAULT_VISUALIZATION_ACTION_NAME } from "@app/lib/api/assistant/actions/names";
@@ -31,6 +35,7 @@ import { renderConversationForModelMultiActions } from "@app/lib/api/assistant/g
 import { getSupportedModelConfig } from "@app/lib/assistant";
 import type { Authenticator } from "@app/lib/auth";
 import { AgentVisualizationAction } from "@app/lib/models/assistant/actions/visualization";
+import { FileResource } from "@app/lib/resources/file_resource";
 import logger from "@app/logger/logger";
 
 interface VisualizationActionBlob {
@@ -108,7 +113,7 @@ export class VisualizationConfigurationServerRunner extends BaseActionConfigurat
       name,
       description:
         description ||
-        "Generates code to represent the requested data in a graph.",
+        "Generates React javascript code that will be run in the browser.",
       inputs: [],
     });
   }
@@ -177,16 +182,61 @@ export class VisualizationConfigurationServerRunner extends BaseActionConfigurat
       }),
     };
 
-    // Turn the conversation into a digest that can be presented to the model.
+    const contentFragmentsText = await Promise.all(
+      conversation.content
+        .flat(1)
+        .filter((m) => isContentFragmentType(m))
+        .map(async (m) => {
+          // This is needed because the ContentFragmentType is not actually narrowed correctly,
+          // this type is actually pretty complex and even if VSCode figure it out, `npm run build` does not.
+          // @flavien knows about this.
+          assert(isContentFragmentType(m));
+          if (!m.fileId) {
+            return;
+          }
+          if (m.contentType.startsWith("text/")) {
+            const file = await FileResource.fetchById(auth, m.fileId);
+            if (file === null) {
+              return null;
+            }
+            const readStream = file.getReadStream({
+              auth,
+              version: "original",
+            });
+            const lines = await readFirstFiveLines(readStream);
+            return lines;
+          } else {
+            return [];
+          }
+        })
+    );
+
+    // The prompt is a list of files that the visualization action can access.
+    // I have empirically found that this format works well for the model to generate the appropriate code,
+    // but feel free to change it if you think otherwise.
+    const prompt =
+      `You have access to the following files:\n` +
+      conversation.content
+        .flat(1)
+        .filter((m) => isContentFragmentType(m))
+        .map((m, i) => {
+          // This is needed because the ContentFragmentType is not actually narrowed correctly,
+          // this type is actually pretty complex and even if VSCode figure it out, `npm run build` does not.
+          // @flavien knows about this.
+          assert(isContentFragmentType(m));
+          return `<file id="${m.fileId}" name="${m.title}" type="${m.contentType}">\n${contentFragmentsText[i]?.join("\n")}(truncated...)</file>`;
+        });
+
     const MIN_GENERATION_TOKENS = 2048;
     const agentModelConfig = getSupportedModelConfig(agentConfiguration.model);
     const modelConversationRes = await renderConversationForModelMultiActions({
       conversation,
       model: agentModelConfig,
-      prompt: "", // There is no prompt for title generation.
+      prompt: prompt,
       allowedTokenCount: agentModelConfig.contextSize - MIN_GENERATION_TOKENS,
       excludeActions: false,
       excludeImages: true,
+      excludeContentFragments: true,
     });
 
     if (modelConversationRes.isErr()) {
@@ -269,6 +319,7 @@ export class VisualizationConfigurationServerRunner extends BaseActionConfigurat
       [
         {
           conversation: renderedConversation,
+          prompt: prompt,
         },
       ],
       {
@@ -381,6 +432,17 @@ export class VisualizationConfigurationServerRunner extends BaseActionConfigurat
       }
     }
 
+    logger.info(
+      {
+        workspaceId: conversation.owner.sId,
+        conversationId: conversation.sId,
+        elapsed: Date.now() - now,
+      },
+      "[ASSISTANT_TRACE] Visualization action execution"
+    );
+
+    await action.update({ runId: await dustRunId, generation });
+
     yield {
       type: "visualization_success",
       created: Date.now(),
@@ -388,24 +450,13 @@ export class VisualizationConfigurationServerRunner extends BaseActionConfigurat
       messageId: agentMessage.sId,
       action: new VisualizationAction({
         id: action.id,
-        agentMessageId: action.agentMessageId,
-        generation,
+        agentMessageId: agentMessage.id,
+        generation: action.generation,
         functionCallId: action.functionCallId,
         functionCallName: action.functionCallName,
         step: action.step,
       }),
     };
-
-    logger.info(
-      {
-        workspaceId: conversation.owner.sId,
-        conversationId: conversation.sId,
-        elapsed: Date.now() - now,
-      },
-      "[ASSISTANT_TRACE] Visualization action execution"
-    );
-
-    await action.update({ runId: await dustRunId, generation });
   }
 }
 
@@ -436,3 +487,31 @@ export async function visualizationActionTypesFromAgentMessageIds(
     });
   });
 }
+
+const readFirstFiveLines = (inputStream: Readable): Promise<string[]> => {
+  return new Promise((resolve, reject) => {
+    const rl: readline.Interface = readline.createInterface({
+      input: inputStream,
+      crlfDelay: Infinity,
+    });
+
+    let lineCount: number = 0;
+    const lines: string[] = [];
+
+    rl.on("line", (line: string) => {
+      lines.push(line);
+      lineCount++;
+      if (lineCount === 5) {
+        rl.close();
+      }
+    });
+
+    rl.on("close", () => {
+      resolve(lines);
+    });
+
+    rl.on("error", (err: Error) => {
+      reject(err);
+    });
+  });
+};
diff --git a/front/lib/api/assistant/generation.ts b/front/lib/api/assistant/generation.ts
index dc46ed204012..30371ba3dd7a 100644
--- a/front/lib/api/assistant/generation.ts
+++ b/front/lib/api/assistant/generation.ts
@@ -45,6 +45,7 @@ export async function renderConversationForModelMultiActions({
   allowedTokenCount,
   excludeActions,
   excludeImages,
+  excludeContentFragments,
 }: {
   conversation: ConversationType;
   model: ModelConfigurationType;
@@ -52,6 +53,7 @@ export async function renderConversationForModelMultiActions({
   allowedTokenCount: number;
   excludeActions?: boolean;
   excludeImages?: boolean;
+  excludeContentFragments?: boolean;
 }): Promise<
   Result<
     {
@@ -198,8 +200,9 @@ export async function renderConversationForModelMultiActions({
       if (res.isErr()) {
         return new Err(res.error);
       }
-
-      messages.push(res.value);
+      if (!excludeContentFragments) {
+        messages.push(res.value);
+      }
     } else {
       assertNever(m);
     }
@@ -449,7 +452,7 @@ export async function constructPromptMultiActions(
     (action) => isVisualizationConfiguration(action)
   );
   if (needVisualizationMetaPrompt) {
-    additionalInstructions += `If mermaid is asked for a graph you can proceed. Otherwise to generate graphs only call the visualization tool. It takes care of writing and rendering the graph above your message.\n`;
+    additionalInstructions += `If mermaid is asked for a graph you can proceed. Otherwise to generate graphs only call the visualization tool. It takes care of writing and rendering the graph above your message.\nNever repeat the generated code to the user.\nIf asked to manipulate CSV files, you can use the tool to generate the code and then use the code in the tool to manipulate the CSV file.\n`;
   }
 
   const providerMetaPrompt = model.metaPrompt;
diff --git a/front/lib/resources/content_fragment_resource.ts b/front/lib/resources/content_fragment_resource.ts
index 599493a45d2a..c79707d6ed0a 100644
--- a/front/lib/resources/content_fragment_resource.ts
+++ b/front/lib/resources/content_fragment_resource.ts
@@ -359,7 +359,7 @@ async function renderFromFileId(
         content: [
           {
             type: "text",
-            text: `<attachment type="${contentType} title="${title}">[Image content interpreted by a vision-enabled model. Description not available in this context.]</attachment>`,
+            text: `<attachment id="${fileId}" type="${contentType} title="${title}">[Image content interpreted by a vision-enabled model. Description not available in this context.]</attachment>`,
           },
         ],
       });
@@ -395,7 +395,7 @@ async function renderFromFileId(
       content: [
         {
           type: "text",
-          text: `<attachment type="${contentType}" title="${title}">\n${content}\n</attachment>`,
+          text: `<attachment id="${fileId}" type="${contentType}" title="${title}">\n${content}\n</attachment>`,
         },
       ],
     });
@@ -440,7 +440,7 @@ export async function renderContentFragmentForModel(
         content: [
           {
             type: "text",
-            text: `<attachment type="${contentType}" title="${title}">\n${content}\n</attachment>`,
+            text: `<attachment id="${fileId}" type="${contentType}" title="${title}">\n${content}\n</attachment>`,
           },
         ],
       });
diff --git a/front/pages/api/w/[wId]/files/[fileId].ts b/front/pages/api/w/[wId]/files/[fileId].ts
index 0e4ed964dd00..41637c51c17c 100644
--- a/front/pages/api/w/[wId]/files/[fileId].ts
+++ b/front/pages/api/w/[wId]/files/[fileId].ts
@@ -81,7 +81,7 @@ async function handler(
             },
           });
         });
-
+        res.setHeader("Content-Type", file.contentType);
         readStream.pipe(res);
         return;
       }
diff --git a/k8s/cloudbuild.yaml b/k8s/cloudbuild.yaml
index 00c1d26805b7..cfb99532b93b 100644
--- a/k8s/cloudbuild.yaml
+++ b/k8s/cloudbuild.yaml
@@ -12,15 +12,18 @@ steps:
       - -f
       - ${_DOCKERFILE_PATH}
       - --build-arg
-      - COMMIT_HASH=${SHORT_SHA}
+      - COMMIT_HASH=$SHORT_SHA
+      - --build-arg
+      - NEXT_PUBLIC_VIZ_URL=https://viz.dust.tt
       - .
-    secretEnv: ["DEPOT_TOKEN"]
+    secretEnv:
+      - "DEPOT_TOKEN"
 
 timeout: 600s
 
 availableSecrets:
   secretManager:
-    - versionName: "projects/$PROJECT_ID/secrets/DEPOT_TOKEN/versions/latest"
+    - versionName: projects/$PROJECT_ID/secrets/DEPOT_TOKEN/versions/latest
       env: DEPOT_TOKEN
 
 options:
diff --git a/types/src/front/assistant/actions/visualization.ts b/types/src/front/assistant/actions/visualization.ts
index 798f697c308c..3ef7796e3d24 100644
--- a/types/src/front/assistant/actions/visualization.ts
+++ b/types/src/front/assistant/actions/visualization.ts
@@ -60,10 +60,125 @@ export function visualizationExtractCode(code: string) {
 }
 
 // This defines the commands that the iframe can send to the host window.
-export type VisualizationRPCCommand = "getCodeToExecute" | "retry" | "getFile";
-export type VisualizationRPCRequest = {
-  command: VisualizationRPCCommand;
-  messageUniqueId: string;
+
+// Common base interface.
+interface VisualizationRPCRequestBase {
   actionId: number;
-  params: unknown;
+  messageUniqueId: string;
+}
+
+// Define parameter types for each command.
+export type GetFileParams = {
+  fileId: string;
 };
+
+export type RetryParams = {
+  errorMessage: string;
+};
+
+// Define a mapped type to extend the base with specific parameters.
+export type VisualizationRPCRequestMap = {
+  getFile: GetFileParams;
+  getCodeToExecute: null;
+  retry: RetryParams;
+};
+
+// Derive the command type from the keys of the request map
+export type VisualizationRPCCommand = keyof VisualizationRPCRequestMap;
+
+// Create a union type for requests based on the mapped type.
+export type VisualizationRPCRequest = {
+  [K in VisualizationRPCCommand]: VisualizationRPCRequestBase & {
+    command: K;
+    params: VisualizationRPCRequestMap[K];
+  };
+}[VisualizationRPCCommand];
+
+export const validCommands: VisualizationRPCCommand[] = [
+  "getFile",
+  "getCodeToExecute",
+  "retry",
+];
+
+// Type guard for getFile.
+export function isGetFileRequest(
+  value: unknown
+): value is VisualizationRPCRequest & {
+  command: "getFile";
+  params: GetFileParams;
+} {
+  if (typeof value !== "object" || value === null) {
+    return false;
+  }
+
+  const v = value as Partial<VisualizationRPCRequest>;
+
+  return (
+    v.command === "getFile" &&
+    typeof v.actionId === "number" &&
+    typeof v.messageUniqueId === "string" &&
+    typeof v.params === "object" &&
+    v.params !== null &&
+    typeof (v.params as GetFileParams).fileId === "string"
+  );
+}
+
+// Type guard for getCodeToExecute.
+export function isGetCodeToExecuteRequest(
+  value: unknown
+): value is VisualizationRPCRequest & {
+  command: "getCodeToExecute";
+  params: null;
+} {
+  if (typeof value !== "object" || value === null) {
+    return false;
+  }
+
+  const v = value as Partial<VisualizationRPCRequest>;
+
+  return (
+    v.command === "getCodeToExecute" &&
+    typeof v.actionId === "number" &&
+    typeof v.messageUniqueId === "string"
+  );
+}
+
+// Type guard for retry.
+export function isRetryRequest(
+  value: unknown
+): value is VisualizationRPCRequest & {
+  command: "retry";
+  params: RetryParams;
+} {
+  if (typeof value !== "object" || value === null) {
+    return false;
+  }
+
+  const v = value as Partial<VisualizationRPCRequest>;
+
+  return (
+    v.command === "retry" &&
+    typeof v.actionId === "number" &&
+    typeof v.messageUniqueId === "string" &&
+    typeof v.params === "object" &&
+    v.params !== null &&
+    typeof (v.params as RetryParams).errorMessage === "string"
+  );
+}
+
+export function isVisualizationRPCRequest(
+  value: unknown
+): value is VisualizationRPCRequest {
+  if (typeof value !== "object" || value === null) {
+    return false;
+  }
+
+  const v = value as Partial<VisualizationRPCRequest>;
+
+  return (
+    typeof v.actionId === "number" &&
+    typeof v.messageUniqueId === "string" &&
+    typeof v.command === "string" &&
+    validCommands.includes(v.command as VisualizationRPCCommand)
+  );
+}
diff --git a/types/src/front/lib/actions/registry.ts b/types/src/front/lib/actions/registry.ts
index 9a75c36f3fb0..8628c1aa444d 100644
--- a/types/src/front/lib/actions/registry.ts
+++ b/types/src/front/lib/actions/registry.ts
@@ -238,7 +238,7 @@ export const DustProdActionRegistry = createActionRegistry({
       workspaceId: PRODUCTION_DUST_APPS_WORKSPACE_ID,
       appId: "tWcuYDj1OE",
       appHash:
-        "a8aaf72e9f5a9f62438fde3a57cb6cfdf461315d6d6e86a03c55248ec5294162",
+        "720fd391bfeb1532f97f4d6ff8bf57f99eef9c56e470c7e3892061de8fb0bd80",
     },
     config: {
       MODEL: {
diff --git a/viz/app/components/VisualizationWrapper.tsx b/viz/app/components/VisualizationWrapper.tsx
index 4b8989121e02..c34f6a97bd70 100644
--- a/viz/app/components/VisualizationWrapper.tsx
+++ b/viz/app/components/VisualizationWrapper.tsx
@@ -2,7 +2,7 @@
 
 import type {
   VisualizationRPCCommand,
-  VisualizationRPCRequest,
+  VisualizationRPCRequestMap,
 } from "@dust-tt/types";
 import { Button, ErrorMessage, Spinner } from "@viz/app/components/Components";
 import * as papaparseAll from "papaparse";
@@ -32,11 +32,11 @@ export function useVisualizationAPI(actionId: number) {
 
   const fetchCode = useCallback(async (): Promise<string | null> => {
     const getCode = makeIframeMessagePassingFunction<
-      { actionId: number },
-      { code?: string }
+      "getCodeToExecute",
+      { code: string }
     >("getCodeToExecute", actionId);
     try {
-      const result = await getCode({ actionId });
+      const result = await getCode(null);
 
       const extractedCode = visualizationExtractCodeNonStreaming(
         result.code ?? ""
@@ -62,7 +62,7 @@ export function useVisualizationAPI(actionId: number) {
   const fetchFile = useCallback(
     async (fileId: string): Promise<File | null> => {
       const getFile = makeIframeMessagePassingFunction<
-        { fileId: string },
+        "getFile",
         { file?: File }
       >("getFile", actionId);
       const res = await getFile({ fileId });
@@ -79,21 +79,23 @@ export function useVisualizationAPI(actionId: number) {
 
   // This retry function sends a command to the host window requesting a retry of a previous
   // operation, typically if the generated code fails.
-  const retry = useCallback(async (): Promise<void> => {
-    const sendRetry = makeIframeMessagePassingFunction("retry", actionId);
-    // TODO(2024-07-24 flav) Pass the error message to the host window.
-    await sendRetry({});
-  }, [actionId]);
+  const retry = useCallback(
+    async (errorMessage: string): Promise<void> => {
+      const sendRetry = makeIframeMessagePassingFunction("retry", actionId);
+      await sendRetry({ errorMessage });
+    },
+    [actionId]
+  );
 
   return { fetchCode, fetchFile, error, retry };
 }
 
 // This function creates a function that sends a command to the host window with templated Input and Output types.
-function makeIframeMessagePassingFunction<Params, Answer>(
-  methodName: VisualizationRPCCommand,
-  actionId: number
-) {
-  return (params?: Params) => {
+function makeIframeMessagePassingFunction<
+  T extends VisualizationRPCCommand,
+  Answer
+>(methodName: T, actionId: number) {
+  return (params: VisualizationRPCRequestMap[T]) => {
     return new Promise<Answer>((resolve, reject) => {
       const messageUniqueId = Math.random().toString();
       const listener = (event: MessageEvent) => {
@@ -113,7 +115,7 @@ function makeIframeMessagePassingFunction<Params, Answer>(
           messageUniqueId,
           actionId,
           params,
-        } satisfies VisualizationRPCRequest,
+        },
         "*"
       );
     });
@@ -179,7 +181,12 @@ export function VisualizationWrapper({ actionId }: { actionId: string }) {
   }, [error]);
 
   if (errored) {
-    return <VisualizationError error={errored} retry={() => retry()} />;
+    return (
+      <VisualizationError
+        error={errored}
+        retry={() => retry(errored.message)}
+      />
+    );
   }
 
   if (!code) {