From 333a557717e6cc21f429a72af8309a5a63c1a77b Mon Sep 17 00:00:00 2001
From: Lucas Massemin <lucasmassemin@gmail.com>
Date: Fri, 15 Nov 2024 14:14:12 +0100
Subject: [PATCH] Can inspect code for visualizations. Layout shift needs fix.
 (#8648)

* Can inspect code for visualizations. Layout shift needs fix.

* ui improvements

---------

Co-authored-by: Lucas <lucas@Dusts-MacBook-Pro.local>
Co-authored-by: Thomas Draier <thomas@dust.tt>
---
 .../actions/VisualizationActionIframe.tsx     | 53 +++++++++++++++++--
 types/src/front/assistant/visualization.ts    | 25 ++++++++-
 viz/app/components/VisualizationWrapper.tsx   | 23 ++++++--
 3 files changed, 93 insertions(+), 8 deletions(-)

diff --git a/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx b/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
index 4b1a07e295f6..ef2f090d4ceb 100644
--- a/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
+++ b/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
@@ -1,6 +1,12 @@
-import { Button, Spinner } from "@dust-tt/sparkle";
-import { MarkdownContentContext } from "@dust-tt/sparkle";
-import { Markdown } from "@dust-tt/sparkle";
+import {
+  Button,
+  CodeBlock,
+  Markdown,
+  MarkdownContentContext,
+  Modal,
+  Page,
+  Spinner,
+} from "@dust-tt/sparkle";
 import type {
   CommandResultMap,
   LightWorkspaceType,
@@ -58,12 +64,14 @@ function useVisualizationDataHandler({
   visualization,
   setContentHeight,
   setErrorMessage,
+  setCodeDrawerOpened,
   vizIframeRef,
   workspaceId,
 }: {
   visualization: Visualization;
   setContentHeight: (v: SetStateAction<number>) => void;
   setErrorMessage: (v: SetStateAction<string | null>) => void;
+  setCodeDrawerOpened: (v: SetStateAction<boolean>) => void;
   vizIframeRef: React.MutableRefObject<HTMLIFrameElement | null>;
   workspaceId: string;
 }) {
@@ -147,6 +155,10 @@ function useVisualizationDataHandler({
           downloadFileFromBlob(data.params.blob, data.params.filename);
           break;
 
+        case "displayCode":
+          setCodeDrawerOpened(true);
+          break;
+
         default:
           assertNever(data);
       }
@@ -160,11 +172,36 @@ function useVisualizationDataHandler({
     getFileBlob,
     setContentHeight,
     setErrorMessage,
+    setCodeDrawerOpened,
     visualization.identifier,
     vizIframeRef,
   ]);
 }
 
+export function CodeDrawer({
+  isOpened,
+  onClose,
+  code,
+}: {
+  isOpened: boolean;
+  onClose: () => void;
+  code: string;
+}) {
+  return (
+    <Modal
+      isOpen={isOpened}
+      onClose={onClose}
+      title="Code for this visualization"
+      variant="side-md"
+      hasChanged={false}
+    >
+      <Page variant="modal">
+        <CodeBlock className="language-jsx">{code}</CodeBlock>
+      </Page>
+    </Modal>
+  );
+}
+
 export function VisualizationActionIframe({
   owner,
   visualization,
@@ -179,7 +216,7 @@ export function VisualizationActionIframe({
   const [contentHeight, setContentHeight] = useState<number>(0);
   const [errorMessage, setErrorMessage] = useState<string | null>(null);
   const [retryClicked, setRetryClicked] = useState(false);
-
+  const [isCodeDrawerOpen, setCodeDrawerOpened] = useState(false);
   const vizIframeRef = useRef<HTMLIFrameElement | null>(null);
 
   const isErrored = !!errorMessage || retryClicked;
@@ -189,6 +226,7 @@ export function VisualizationActionIframe({
     workspaceId: owner.sId,
     setContentHeight,
     setErrorMessage,
+    setCodeDrawerOpened,
     vizIframeRef,
   });
 
@@ -226,6 +264,13 @@ export function VisualizationActionIframe({
           <Spinner size="xl" />
         </div>
       )}
+      {code && (
+        <CodeDrawer
+          isOpened={isCodeDrawerOpen}
+          onClose={() => setCodeDrawerOpened(false)}
+          code={code}
+        />
+      )}
       <div
         className={classNames(
           "relative w-full overflow-hidden",
diff --git a/types/src/front/assistant/visualization.ts b/types/src/front/assistant/visualization.ts
index 5ae2ab0065d0..1f0dd57108ed 100644
--- a/types/src/front/assistant/visualization.ts
+++ b/types/src/front/assistant/visualization.ts
@@ -32,6 +32,7 @@ export type VisualizationRPCRequestMap = {
   setContentHeight: SetContentHeightParams;
   setErrorMessage: setErrorMessageParams;
   downloadFileRequest: DownloadFileRequestParams;
+  displayCode: null;
 };
 
 // Derive the command type from the keys of the request map
@@ -60,6 +61,7 @@ export interface CommandResultMap {
   downloadFileRequest: { blob: Blob; filename?: string };
   setContentHeight: void;
   setErrorMessage: void;
+  displayCode: void;
 }
 
 // TODO(@fontanierh): refactor all these guards to use io-ts instead of manual checks.
@@ -170,6 +172,26 @@ export function isDownloadFileRequest(
   );
 }
 
+// Type guard for getCodeToExecute.
+export function isDisplayCodeRequest(
+  value: unknown
+): value is VisualizationRPCRequest & {
+  command: "displayCode";
+  params: null;
+} {
+  if (typeof value !== "object" || value === null) {
+    return false;
+  }
+
+  const v = value as Partial<VisualizationRPCRequest>;
+
+  return (
+    v.command === "displayCode" &&
+    typeof v.identifier === "string" &&
+    typeof v.messageUniqueId === "string"
+  );
+}
+
 export function isVisualizationRPCRequest(
   value: unknown
 ): value is VisualizationRPCRequest {
@@ -182,6 +204,7 @@ export function isVisualizationRPCRequest(
     isGetFileRequest(value) ||
     isDownloadFileRequest(value) ||
     isSetContentHeightRequest(value) ||
-    isSetErrorMessageRequest(value)
+    isSetErrorMessageRequest(value) ||
+    isDisplayCodeRequest(value)
   );
 }
diff --git a/viz/app/components/VisualizationWrapper.tsx b/viz/app/components/VisualizationWrapper.tsx
index 57f6d5a04412..145d447c842f 100644
--- a/viz/app/components/VisualizationWrapper.tsx
+++ b/viz/app/components/VisualizationWrapper.tsx
@@ -14,7 +14,7 @@ import { importCode, Runner } from "react-runner";
 import * as rechartsAll from "recharts";
 import { useResizeDetector } from "react-resize-detector";
 import { ErrorBoundary } from "@viz/app/components/ErrorBoundary";
-import { Download } from "lucide-react";
+import { Download, SquareTerminal } from "lucide-react";
 import { toBlob } from "html-to-image";
 
 export function useVisualizationAPI(
@@ -83,12 +83,20 @@ export function useVisualizationAPI(
     [sendCrossDocumentMessage]
   );
 
+  const displayCode = useCallback(
+    async () => {
+      await sendCrossDocumentMessage("displayCode", null);
+    },
+    [sendCrossDocumentMessage]
+  );
+
   return {
     error,
     fetchCode,
     fetchFile,
     sendHeightToParent,
     downloadFile,
+    displayCode,
   };
 }
 
@@ -182,7 +190,7 @@ export function VisualizationWrapper({
 
   const [errored, setErrorMessage] = useState<Error | null>(null);
 
-  const { fetchCode, fetchFile, error, sendHeightToParent, downloadFile } = api;
+  const { fetchCode, fetchFile, error, sendHeightToParent, downloadFile, displayCode } = api;
 
   const memoizedDownloadFile = useDownloadFileCallback(downloadFile);
 
@@ -249,6 +257,10 @@ export function VisualizationWrapper({
     }
   }, [ref, downloadFile]);
 
+    const handleDisplayCode = useCallback(async () => {
+    await displayCode();
+  }, [displayCode]);
+
   useEffect(() => {
     if (error) {
       setErrorMessage(error);
@@ -266,12 +278,17 @@ export function VisualizationWrapper({
 
   return (
     <div className="relative group/viz">
+      <div className="flex flex-row gap-2 absolute top-2 right-2 bg-white rounded transition opacity-0 group-hover/viz:opacity-100 z-50">
       <button
         onClick={handleScreenshotDownload}
-        className="absolute top-2 right-2 bg-white p-2 rounded shadow hover:bg-gray-100 transition opacity-0 group-hover/viz:opacity-100 z-50"
+        className="hover:bg-slate-200 rounded p-2 border border-slate-200"
       >
         <Download size={20} />
       </button>
+      <button className="hover:bg-slate-200 rounded p-2 border border-slate-200" onClick={handleDisplayCode}>
+        <SquareTerminal size={20} />
+      </button>
+      </div>
       <div ref={ref}>
         <Runner
           code={runnerParams.code}