enh: remove viz action and replace by viz content and system prompt

dust-tt · Aug 1, 2024 · e12e68d · e12e68d
1 parent 5752de8
commit e12e68d
Show file tree

Hide file tree

Showing 41 changed files with 882 additions and 1,286 deletions.
diff --git a/front/components/actions/types.ts b/front/components/actions/types.ts
@@ -5,7 +5,6 @@ import { DustAppRunActionDetails } from "@app/components/actions/dust_app_run/Du
 import { ProcessActionDetails } from "@app/components/actions/process/ProcessActionDetails";
 import { RetrievalActionDetails } from "@app/components/actions/retrieval/RetrievalActionDetails";
 import { TablesQueryActionDetails } from "@app/components/actions/tables_query/TablesQueryActionDetails";
-import { VisualizationActionDetails } from "@app/components/actions/visualization/VisualizationActionDetails";
 import { WebsearchActionDetails } from "@app/components/actions/websearch/WebsearchActionDetails";
 
 export interface ActionDetailsComponentBaseProps<
@@ -51,10 +50,6 @@ const actionsSpecification: ActionSpecifications = {
     detailsComponent: BrowseActionDetails,
     runningLabel: "Browsing page",
   },
-  visualization_action: {
-    detailsComponent: VisualizationActionDetails,
-    runningLabel: "Analyzing request",
-  },
 };
 
 export function getActionSpecification<T extends ActionType>(

diff --git a/front/components/actions/visualization/VisualizationActionDetails.tsx b/front/components/actions/visualization/VisualizationActionDetails.tsx
diff --git a/front/components/assistant/AssistantActions.tsx b/front/components/assistant/AssistantActions.tsx
@@ -199,6 +199,7 @@ export function RemoveAssistantFromWorkspaceDialog({
             actions: detailedConfiguration.actions,
             templateId: agentConfiguration.templateId,
             maxStepsPerRun: agentConfiguration.maxStepsPerRun,
+            visualizationEnabled: agentConfiguration.visualizationEnabled,
           },
         };
 

diff --git a/front/components/assistant/AssistantDetails.tsx b/front/components/assistant/AssistantDetails.tsx
@@ -10,7 +10,6 @@ import {
   Page,
   PlanetIcon,
   ServerIcon,
-  ShapesIcon,
   Spinner,
   Tree,
 } from "@dust-tt/sparkle";
@@ -34,7 +33,6 @@ import {
   isProcessConfiguration,
   isRetrievalConfiguration,
   isTablesQueryConfiguration,
-  isVisualizationConfiguration,
   isWebsearchConfiguration,
 } from "@dust-tt/types";
 import { useCallback, useContext, useEffect, useMemo, useState } from "react";
@@ -315,19 +313,6 @@ export function AssistantDetails({
               </div>
             ) : isBrowseConfiguration(action) ? (
               false
-            ) : isVisualizationConfiguration(action) ? (
-              <div className="flex flex-col gap-2" key={`action-${index}`}>
-                <div className="text-lg font-bold text-element-800">
-                  Visualization
-                </div>
-                <div className="flex items-center gap-2">
-                  <Icon visual={ShapesIcon} size="xs" />
-                  <div>
-                    Assistant can generate graphs to visually represent your
-                    data.
-                  </div>
-                </div>
-              </div>
             ) : (
               !isRetrievalConfiguration(action) && assertNever(action)
             )

diff --git a/front/components/assistant/conversation/AgentMessage.tsx b/front/components/assistant/conversation/AgentMessage.tsx
@@ -33,11 +33,9 @@ import type {
 import {
   assertNever,
   isRetrievalActionType,
-  isVisualizationActionType,
   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";
@@ -88,8 +86,8 @@ export function AgentMessage({
   const [streamedAgentMessage, setStreamedAgentMessage] =
     useState<AgentMessageType>(message);
 
-  const [streamedVisualizations, setStreamedVisualizations] = useState<
-    { actionId: number; visualization: string }[]
+  const [visualizations, setVisualizations] = useState<
+    { visualization: string; complete: boolean }[]
   >([]);
 
   const [isRetryHandlerProcessing, setIsRetryHandlerProcessing] =
@@ -103,6 +101,17 @@ export function AgentMessage({
     { index: number; document: RetrievalDocumentType | WebsearchResultType }[]
   >([]);
 
+  useEffect(() => {
+    if (message.status === "succeeded") {
+      setVisualizations(
+        message.visualizations.map((v) => ({
+          visualization: v,
+          complete: true,
+        }))
+      );
+    }
+  }, [message.status, message.visualizations]);
+
   const shouldStream = (() => {
     if (message.status !== "created") {
       return false;
@@ -180,7 +189,6 @@ export function AgentMessage({
       case "process_params":
       case "websearch_params":
       case "browse_params":
-      case "visualization_params":
         setStreamedAgentMessage((m) => {
           return updateMessageWithAction(m, event.action);
         });
@@ -203,13 +211,27 @@ export function AgentMessage({
             ...event.message,
           };
         });
-        setStreamedVisualizations([]);
         break;
       }
 
       case "generation_tokens": {
         switch (event.classification) {
           case "closing_delimiter":
+            if (event.delimiterClassification === "visualization") {
+              // If we receive a closing delimiter for a visualization, we can
+              // consider the last viz to be complete.
+              setVisualizations((v) => {
+                const lastViz = v[v.length - 1];
+                if (lastViz) {
+                  return [
+                    ...v.slice(0, v.length - 1),
+                    { ...lastViz, complete: true },
+                  ];
+                }
+                return v;
+              });
+            }
+            break;
           case "opening_delimiter":
             break;
           case "tokens":
@@ -229,30 +251,29 @@ export function AgentMessage({
               };
             });
             break;
+          case "visualization":
+            // Append new content to the last viz, or create a new one if there
+            // is none or the last one is complete.
+            setVisualizations((v) => {
+              const lastViz = v[v.length - 1];
+              if (lastViz && !lastViz.complete) {
+                return [
+                  ...v.slice(0, v.length - 1),
+                  {
+                    visualization: lastViz.visualization + event.text,
+                    complete: false,
+                  },
+                ];
+              }
+              return [...v, { visualization: event.text, complete: false }];
+            });
+            break;
           default:
-            assertNever(event.classification);
+            assertNever(event);
         }
         break;
       }
 
-      case "visualization_generation_tokens":
-        setStreamedVisualizations((m) => {
-          const actionId = event.actionId;
-          const tokens = event.text;
-          const index = m.findIndex((v) => v.actionId === actionId);
-          if (index === -1) {
-            return [...m, { actionId, visualization: tokens }];
-          } else {
-            return m.map((v) => {
-              if (v.actionId === actionId) {
-                return { ...v, visualization: v.visualization + tokens };
-              }
-              return v;
-            });
-          }
-        });
-        break;
-
       default:
         assertNever(event);
     }
@@ -469,7 +490,7 @@ export function AgentMessage({
           references: references,
           streaming: shouldStream,
           lastTokenClassification: lastTokenClassification,
-          streamedVisualizations,
+          visualizations,
         })}
       </div>
       {/* Invisible div to act as a scroll anchor for detecting when the user has scrolled to the bottom */}
@@ -482,13 +503,13 @@ export function AgentMessage({
     references,
     streaming,
     lastTokenClassification,
-    streamedVisualizations,
+    visualizations,
   }: {
     agentMessage: AgentMessageType;
     references: { [key: string]: RetrievalDocumentType | WebsearchResultType };
     streaming: boolean;
     lastTokenClassification: null | "tokens" | "chain_of_thought";
-    streamedVisualizations: { actionId: number; visualization: string }[];
+    visualizations: { visualization: string; complete: boolean }[];
   }) {
     if (agentMessage.status === "failed") {
       return (
@@ -510,25 +531,16 @@ export function AgentMessage({
       <div className="flex flex-col gap-y-4">
         <AgentMessageActions agentMessage={agentMessage} size={size} />
         <>
-          {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}
-                />
-              );
-            })}
+          {visualizations.map((v, i) => {
+            return (
+              <VisualizationActionIframe
+                visualization={{ ...v, index: i }}
+                key={i}
+                onRetry={() => retryHandler(agentMessage)}
+                owner={owner}
+              />
+            );
+          })}
         </>
 
         {agentMessage.chainOfThought?.length ? (