From bc469a0a1886cb62ffda5f5e075345e3bd6b55e4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daphn=C3=A9=20Popin?= <PopDaph@users.noreply.github.com>
Date: Wed, 20 Dec 2023 15:38:17 +0100
Subject: [PATCH] Agent DbQuery action: improve display of streamed events
 (#2973)

---
 .../assistant/conversation/AgentMessage.tsx   |   1 +
 .../conversation/DatabaseQueryAction.tsx      | 139 ++++++++++++------
 .../api/assistant/actions/database_query.ts   |  34 ++++-
 front/lib/api/assistant/agent.ts              |   1 +
 front/lib/api/assistant/conversation.ts       |   2 +
 front/lib/api/assistant/pubsub.ts             |   2 +
 .../api/assistant/actions/database_query.ts   |   8 +
 types/src/front/lib/api/assistant/agent.ts    |   8 +-
 8 files changed, 144 insertions(+), 51 deletions(-)

diff --git a/front/components/assistant/conversation/AgentMessage.tsx b/front/components/assistant/conversation/AgentMessage.tsx
index c436eb9533d2..66e8f36ef9e7 100644
--- a/front/components/assistant/conversation/AgentMessage.tsx
+++ b/front/components/assistant/conversation/AgentMessage.tsx
@@ -128,6 +128,7 @@ export function AgentMessage({
       case "dust_app_run_params":
       case "dust_app_run_block":
       case "database_query_params":
+      case "database_query_output":
         setStreamedAgentMessage((m) => {
           return { ...m, action: event.action };
         });
diff --git a/front/components/assistant/conversation/DatabaseQueryAction.tsx b/front/components/assistant/conversation/DatabaseQueryAction.tsx
index 8c814bcad87b..f5ca39f8fa61 100644
--- a/front/components/assistant/conversation/DatabaseQueryAction.tsx
+++ b/front/components/assistant/conversation/DatabaseQueryAction.tsx
@@ -4,6 +4,7 @@ import {
   Chip,
   Icon,
   Spinner,
+  Tooltip,
 } from "@dust-tt/sparkle";
 import { DatabaseQueryActionType } from "@dust-tt/types";
 import dynamic from "next/dynamic";
@@ -20,72 +21,114 @@ export default function DatabaseQueryAction({
 }: {
   databaseQueryAction: DatabaseQueryActionType;
 }) {
-  const [outputVisible, setOutputVisible] = useState(false);
+  const [isOutputExpanded, setIsOutputExpanded] = useState(false);
 
+  // Extracting question from the params
+  const params = databaseQueryAction.params;
+  const question =
+    typeof params?.question === "string" ? params.question : null;
+
+  // Extracting query and result from the output
   const output = databaseQueryAction.output;
-  const query = output?.query;
+  const query = typeof output?.query === "string" ? output.query : null;
+  const noQuery = output?.no_query === true;
+  const results = output?.results;
+
+  const isQueryStepCompleted = noQuery || query;
+  const isOutputStepCompleted = noQuery || (query && results);
+
+  const trimText = (text: string, maxLength = 20) => {
+    const t = text.replaceAll("\n", " ");
+    return t.length > maxLength ? t.substring(0, maxLength) + "..." : t;
+  };
 
   return (
     <>
-      <div className="grid grid-cols-[auto,1fr] gap-2">
-        <div className="grid-cols-auto grid items-center">
-          {!output ? (
-            <div>
-              <Spinner size="sm" />
-            </div>
-          ) : (
+      {question && (
+        <div className="flex flex-row items-center gap-2 pb-2">
+          <div className="flex flex-col items-start text-xs font-bold text-element-600">
+            <div className="flex">Question:</div>
+          </div>
+          <Tooltip label={question}>
+            <Chip color="slate" label={trimText(question)} />
+          </Tooltip>
+        </div>
+      )}
+
+      {!isQueryStepCompleted && (
+        <div>
+          <div className="pb-2 text-xs font-bold text-element-600">
+            Generating query...
+          </div>
+          <Spinner size="sm" />
+        </div>
+      )}
+
+      {isQueryStepCompleted && (
+        <div className="grid grid-cols-[auto,1fr] gap-2 pb-2">
+          <div className="grid-cols-auto grid items-center">
             <div className="text-xs font-bold text-element-600">
-              <span>{query ? "Query Executed:" : "Result: "}</span>
+              <span>Query:</span>
             </div>
-          )}
-        </div>
-        {!!output && (
+          </div>
           <div className="row-span-1 select-none">
             <div
               className="cursor-pointer"
               onClick={() => {
-                setOutputVisible(!outputVisible);
+                setIsOutputExpanded(!isOutputExpanded);
               }}
             >
               <Chip color="purple">
-                {query ? query : "No query generated, expand to see why"}
-                <Icon
-                  visual={outputVisible ? ChevronDownIcon : ChevronRightIcon}
-                  size="xs"
-                />
+                {query ? query : "No query generated"}
+                {(noQuery || results) && (
+                  <Icon
+                    visual={
+                      isOutputExpanded ? ChevronDownIcon : ChevronRightIcon
+                    }
+                    size="xs"
+                  />
+                )}
               </Chip>
             </div>
           </div>
-        )}
-        {outputVisible && (
-          <div className="col-start-2 row-span-1 max-h-48 overflow-auto rounded-md bg-structure-100">
-            <SyntaxHighlighter
-              className="h-full w-full rounded-md text-xs"
-              style={{
-                "hljs-number": {
-                  color: amber["500"],
-                },
-                "hljs-literal": {
-                  color: amber["500"],
-                },
-                "hljs-string": {
-                  color: emerald["600"],
-                },
-                hljs: {
-                  display: "block",
-                  overflowX: "auto",
-                  color: slate["700"],
-                  padding: "1em",
-                },
-              }}
-              language={"json"}
-              PreTag="div"
-            >
-              {JSON.stringify(output, null, 2)}
-            </SyntaxHighlighter>
+          {isOutputExpanded && (
+            <div className="col-start-2 row-span-1 max-h-48 overflow-auto rounded-md bg-structure-100">
+              <SyntaxHighlighter
+                className="h-full w-full rounded-md text-xs"
+                style={{
+                  "hljs-number": {
+                    color: amber["500"],
+                  },
+                  "hljs-literal": {
+                    color: amber["500"],
+                  },
+                  "hljs-string": {
+                    color: emerald["600"],
+                  },
+                  hljs: {
+                    display: "block",
+                    overflowX: "auto",
+                    color: slate["700"],
+                    padding: "1em",
+                  },
+                }}
+                language={"json"}
+                PreTag="div"
+              >
+                {JSON.stringify(output, null, 2)}
+              </SyntaxHighlighter>
+            </div>
+          )}
+        </div>
+      )}
+      {isQueryStepCompleted && !isOutputStepCompleted && (
+        <div>
+          <div className="pb-2 text-xs font-bold text-element-600">
+            Running query...
           </div>
-        )}
-      </div>
+          <Spinner size="sm" />
+        </div>
+      )}
     </>
   );
 }
diff --git a/front/lib/api/assistant/actions/database_query.ts b/front/lib/api/assistant/actions/database_query.ts
index ca0bd9378c93..731fb14144fb 100644
--- a/front/lib/api/assistant/actions/database_query.ts
+++ b/front/lib/api/assistant/actions/database_query.ts
@@ -5,6 +5,7 @@ import {
   ConversationType,
   DatabaseQueryActionType,
   DatabaseQueryErrorEvent,
+  DatabaseQueryOutputEvent,
   DatabaseQueryParamsEvent,
   DatabaseQuerySuccessEvent,
   DustProdActionRegistry,
@@ -119,7 +120,10 @@ export async function* runDatabaseQuery({
   userMessage: UserMessageType;
   agentMessage: AgentMessageType;
 }): AsyncGenerator<
-  DatabaseQueryErrorEvent | DatabaseQuerySuccessEvent | DatabaseQueryParamsEvent
+  | DatabaseQueryErrorEvent
+  | DatabaseQuerySuccessEvent
+  | DatabaseQueryParamsEvent
+  | DatabaseQueryOutputEvent
 > {
   // Checking authorizations
   const owner = auth.workspace();
@@ -282,8 +286,36 @@ export async function* runDatabaseQuery({
         return;
       }
 
+      if (event.content.block_name === "SQL") {
+        let tmpOutput = null;
+        if (e.value) {
+          const sql = e.value as string;
+          tmpOutput = { query: sql };
+        } else {
+          tmpOutput = { no_query: true };
+        }
+        yield {
+          type: "database_query_output",
+          created: Date.now(),
+          configurationId: configuration.sId,
+          messageId: agentMessage.sId,
+          action: {
+            id: action.id,
+            type: "database_query_action",
+            dataSourceWorkspaceId: action.dataSourceWorkspaceId,
+            dataSourceId: action.dataSourceId,
+            databaseId: action.databaseId,
+            params: action.params,
+            output: tmpOutput,
+          },
+        };
+      }
+
       if (event.content.block_name === "OUTPUT" && e.value) {
         output = JSON.parse(e.value as string);
+        if (!output.query) {
+          output.no_query = true;
+        }
       }
     }
   }
diff --git a/front/lib/api/assistant/agent.ts b/front/lib/api/assistant/agent.ts
index 7644e918a85e..1b4aa1ee8598 100644
--- a/front/lib/api/assistant/agent.ts
+++ b/front/lib/api/assistant/agent.ts
@@ -302,6 +302,7 @@ export async function* runAgent(
       for await (const event of eventStream) {
         switch (event.type) {
           case "database_query_params":
+          case "database_query_output":
             yield event;
             break;
           case "database_query_error":
diff --git a/front/lib/api/assistant/conversation.ts b/front/lib/api/assistant/conversation.ts
index eb49b25443e5..bdfd875e0545 100644
--- a/front/lib/api/assistant/conversation.ts
+++ b/front/lib/api/assistant/conversation.ts
@@ -359,6 +359,7 @@ async function batchRenderAgentMessages(
           dataSourceWorkspaceId: action.dataSourceWorkspaceId,
           dataSourceId: action.dataSourceId,
           databaseId: action.databaseId,
+          params: action.params,
           output: action.output,
         };
       });
@@ -1956,6 +1957,7 @@ async function* streamRunAgentEvents(
       case "dust_app_run_params":
       case "dust_app_run_block":
       case "database_query_params":
+      case "database_query_output":
         yield event;
         break;
       case "generation_tokens":
diff --git a/front/lib/api/assistant/pubsub.ts b/front/lib/api/assistant/pubsub.ts
index 63119c89c5d9..56f6470be4c1 100644
--- a/front/lib/api/assistant/pubsub.ts
+++ b/front/lib/api/assistant/pubsub.ts
@@ -158,6 +158,7 @@ async function handleUserMessageEvents(
               case "dust_app_run_params":
               case "dust_app_run_block":
               case "database_query_params":
+              case "database_query_output":
               case "agent_error":
               case "agent_action_success":
               case "generation_tokens":
@@ -282,6 +283,7 @@ export async function retryAgentMessageWithPubSub(
               case "dust_app_run_params":
               case "dust_app_run_block":
               case "database_query_params":
+              case "database_query_output":
               case "agent_error":
               case "agent_action_success":
               case "generation_tokens":
diff --git a/types/src/front/lib/api/assistant/actions/database_query.ts b/types/src/front/lib/api/assistant/actions/database_query.ts
index de3033ef8b1a..3586f7f7a891 100644
--- a/types/src/front/lib/api/assistant/actions/database_query.ts
+++ b/types/src/front/lib/api/assistant/actions/database_query.ts
@@ -26,3 +26,11 @@ export type DatabaseQueryParamsEvent = {
   messageId: string;
   action: DatabaseQueryActionType;
 };
+
+export type DatabaseQueryOutputEvent = {
+  type: "database_query_output";
+  created: number;
+  configurationId: string;
+  messageId: string;
+  action: DatabaseQueryActionType;
+};
diff --git a/types/src/front/lib/api/assistant/agent.ts b/types/src/front/lib/api/assistant/agent.ts
index 7660c80dd644..93528860765b 100644
--- a/types/src/front/lib/api/assistant/agent.ts
+++ b/types/src/front/lib/api/assistant/agent.ts
@@ -6,7 +6,10 @@ import {
   AgentActionType,
   AgentMessageType,
 } from "../../../../front/assistant/conversation";
-import { DatabaseQueryParamsEvent } from "../../../../front/lib/api/assistant/actions/database_query";
+import {
+  DatabaseQueryOutputEvent,
+  DatabaseQueryParamsEvent,
+} from "../../../../front/lib/api/assistant/actions/database_query";
 import {
   DustAppRunBlockEvent,
   DustAppRunParamsEvent,
@@ -41,7 +44,8 @@ export type AgentActionEvent =
   | RetrievalParamsEvent
   | DustAppRunParamsEvent
   | DustAppRunBlockEvent
-  | DatabaseQueryParamsEvent;
+  | DatabaseQueryParamsEvent
+  | DatabaseQueryOutputEvent;
 
 // Event sent once the action is completed, we're moving to generating a message if applicable.
 export type AgentActionSuccessEvent = {