From faf7bdff359c5a5a3a8a2715dfd79ba357f852d4 Mon Sep 17 00:00:00 2001
From: Henry Fontanier <h.fontanier@gmail.com>
Date: Thu, 1 Aug 2024 15:35:08 +0200
Subject: [PATCH] feat: add ability to toggle viz from builder (gated) (#6624)

Co-authored-by: Henry Fontanier <henry@dust.tt>
---
 front/components/assistant/TryAssistant.tsx   |  2 ++
 .../assistant_builder/ActionsScreen.tsx       | 26 +++++++++++++++++++
 .../assistant_builder/AssistantBuilder.tsx    |  1 +
 .../submitAssistantBuilderForm.ts             |  3 +--
 front/components/assistant_builder/types.ts   |  3 +++
 front/lib/api/assistant/templates.ts          |  1 +
 .../[wId]/builder/assistants/[aId]/index.tsx  |  1 +
 .../pages/w/[wId]/builder/assistants/new.tsx  |  1 +
 types/src/front/assistant/agent.ts            |  1 +
 9 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/front/components/assistant/TryAssistant.tsx b/front/components/assistant/TryAssistant.tsx
index 46d145215e33..a1b0b99604ee 100644
--- a/front/components/assistant/TryAssistant.tsx
+++ b/front/components/assistant/TryAssistant.tsx
@@ -93,6 +93,7 @@ export function usePreviewAssistant({
         generationSettings: builderState.generationSettings,
         actions: builderState.actions,
         maxStepsPerRun: builderState.maxStepsPerRun,
+        visualizationEnabled: builderState.visualizationEnabled,
         templateId: builderState.templateId,
       },
       agentConfigurationId: null,
@@ -131,6 +132,7 @@ export function usePreviewAssistant({
     builderState.actions,
     builderState.maxStepsPerRun,
     builderState.templateId,
+    builderState.visualizationEnabled,
     sendNotification,
   ]);
 
diff --git a/front/components/assistant_builder/ActionsScreen.tsx b/front/components/assistant_builder/ActionsScreen.tsx
index 634f08ad614b..6591caa14428 100644
--- a/front/components/assistant_builder/ActionsScreen.tsx
+++ b/front/components/assistant_builder/ActionsScreen.tsx
@@ -2,6 +2,7 @@ import {
   BookOpenIcon,
   Button,
   CardButton,
+  Checkbox,
   ContentMessage,
   DropdownMenu,
   Icon,
@@ -346,6 +347,31 @@ export default function ActionsScreen({
           </div>
         </div>
       </div>
+      <div className="flex flex-col gap-8">
+        {owner.flags.includes("visualization_action_flag") && (
+          <div className="flex flex-row gap-2">
+            <Checkbox
+              checked={builderState.visualizationEnabled}
+              onChange={() => {
+                setEdited(true);
+                setBuilderState((state) => ({
+                  ...state,
+                  visualizationEnabled: !state.visualizationEnabled,
+                }));
+              }}
+              variant="checkable"
+            />
+            <div>
+              <div className="flex w-full gap-1 font-medium text-element-900">
+                Visualization
+              </div>
+              <div className="w-full truncate text-base text-element-700">
+                The assistant can create graphs and charts
+              </div>
+            </div>
+          </div>
+        )}
+      </div>
     </>
   );
 }
diff --git a/front/components/assistant_builder/AssistantBuilder.tsx b/front/components/assistant_builder/AssistantBuilder.tsx
index 879fc3188e78..3dbd7e1739f6 100644
--- a/front/components/assistant_builder/AssistantBuilder.tsx
+++ b/front/components/assistant_builder/AssistantBuilder.tsx
@@ -115,6 +115,7 @@ export default function AssistantBuilder({
           maxStepsPerRun:
             initialBuilderState.maxStepsPerRun ??
             getDefaultAssistantState().maxStepsPerRun,
+          visualizationEnabled: initialBuilderState.visualizationEnabled,
           templateId: initialBuilderState.templateId,
         }
       : {
diff --git a/front/components/assistant_builder/submitAssistantBuilderForm.ts b/front/components/assistant_builder/submitAssistantBuilderForm.ts
index 113d30c0ba77..772f75ded2d4 100644
--- a/front/components/assistant_builder/submitAssistantBuilderForm.ts
+++ b/front/components/assistant_builder/submitAssistantBuilderForm.ts
@@ -200,8 +200,7 @@ export async function submitAssistantBuilderForm({
           temperature: builderState.generationSettings.temperature,
         },
         maxStepsPerRun,
-        // TODO(@fontanierh): support viz in the builder
-        visualizationEnabled: false,
+        visualizationEnabled: builderState.visualizationEnabled,
         templateId: builderState.templateId,
       },
     };
diff --git a/front/components/assistant_builder/types.ts b/front/components/assistant_builder/types.ts
index 7b13a996f588..fb1bc70ca5e8 100644
--- a/front/components/assistant_builder/types.ts
+++ b/front/components/assistant_builder/types.ts
@@ -169,6 +169,7 @@ export type AssistantBuilderState = {
   };
   actions: Array<AssistantBuilderActionConfiguration>;
   maxStepsPerRun: number | null;
+  visualizationEnabled: boolean;
   templateId: string | null;
 };
 
@@ -184,6 +185,7 @@ export type AssistantBuilderInitialState = {
   } | null;
   actions: Array<AssistantBuilderActionConfiguration>;
   maxStepsPerRun: number | null;
+  visualizationEnabled: boolean;
   templateId: string | null;
 };
 
@@ -204,6 +206,7 @@ export function getDefaultAssistantState() {
       temperature: 0.7,
     },
     maxStepsPerRun: 3,
+    visualizationEnabled: false,
     templateId: null,
   } satisfies AssistantBuilderState;
 }
diff --git a/front/lib/api/assistant/templates.ts b/front/lib/api/assistant/templates.ts
index 2497870747e9..0e8a0c7dcf82 100644
--- a/front/lib/api/assistant/templates.ts
+++ b/front/lib/api/assistant/templates.ts
@@ -42,6 +42,7 @@ export async function generateMockAgentConfigurationFromTemplate(
     name: template.handle,
     scope: flow === "personal_assistants" ? "private" : "workspace",
     pictureUrl: template.pictureUrl,
+    visualizationEnabled: false,
     isTemplate: true,
   });
 }
diff --git a/front/pages/w/[wId]/builder/assistants/[aId]/index.tsx b/front/pages/w/[wId]/builder/assistants/[aId]/index.tsx
index 4c3269fe9957..883476ad8618 100644
--- a/front/pages/w/[wId]/builder/assistants/[aId]/index.tsx
+++ b/front/pages/w/[wId]/builder/assistants/[aId]/index.tsx
@@ -174,6 +174,7 @@ export default function EditAssistant({
           temperature: agentConfiguration.model.temperature,
         },
         actions,
+        visualizationEnabled: agentConfiguration.visualizationEnabled,
         maxStepsPerRun: agentConfiguration.maxStepsPerRun,
         templateId: agentConfiguration.templateId,
       }}
diff --git a/front/pages/w/[wId]/builder/assistants/new.tsx b/front/pages/w/[wId]/builder/assistants/new.tsx
index 00431f8edec2..83f8cb7ce0cf 100644
--- a/front/pages/w/[wId]/builder/assistants/new.tsx
+++ b/front/pages/w/[wId]/builder/assistants/new.tsx
@@ -223,6 +223,7 @@ export default function CreateAssistant({
                 temperature: agentConfiguration.model.temperature,
               },
               maxStepsPerRun: agentConfiguration.maxStepsPerRun ?? null,
+              visualizationEnabled: agentConfiguration.visualizationEnabled,
               templateId: templateId,
             }
           : null
diff --git a/types/src/front/assistant/agent.ts b/types/src/front/assistant/agent.ts
index 20f14661f208..8c9217982004 100644
--- a/types/src/front/assistant/agent.ts
+++ b/types/src/front/assistant/agent.ts
@@ -196,6 +196,7 @@ export interface TemplateAgentConfigurationType {
   instructions: string | null;
   isTemplate: true;
   maxStepsPerRun?: number;
+  visualizationEnabled: boolean;
 }
 
 export const MAX_STEPS_USE_PER_RUN_LIMIT = 8;