refactor: tables query action (#3147)

* wip refactor tables query action

* build works

* Refacto: databases (front) part 2

* tmp

* fix action rendering

* fix details

* clean

* select several tables

* hide already selected

* cleanup yield

* concurrently fetch tables + use join for datasources

* nit

* fix some linting violations

* clarify description

* index file

* DustAppParameters

* nit

* fix reb

* lint

---------

Co-authored-by: Henry Fontanier <henry@dust.tt>

front/admin/db.ts

@@ -1,7 +1,5 @@
 import {
   AgentConfiguration,
-  AgentDatabaseQueryAction,
-  AgentDatabaseQueryConfiguration,
   AgentDataSourceConfiguration,
   AgentDustAppRunAction,
   AgentDustAppRunConfiguration,
@@ -39,6 +37,11 @@ import {
   XP1Run,
   XP1User,
 } from "@app/lib/models";
+import {
+  AgentTablesQueryAction,
+  AgentTablesQueryConfiguration,
+  AgentTablesQueryConfigurationTable,
+} from "@app/lib/models/assistant/actions/tables_query";
 import {
   AgentUserRelation,
   GlobalAgentSettings,
@@ -71,8 +74,9 @@ async function main() {
 
   await AgentDustAppRunConfiguration.sync({ alter: true });
   await AgentDustAppRunAction.sync({ alter: true });
-  await AgentDatabaseQueryConfiguration.sync({ alter: true });
-  await AgentDatabaseQueryAction.sync({ alter: true });
+  await AgentTablesQueryConfiguration.sync({ alter: true });
+  await AgentTablesQueryConfigurationTable.sync({ alter: true });
+  await AgentTablesQueryAction.sync({ alter: true });
 
   await AgentGenerationConfiguration.sync({ alter: true });
   await AgentRetrievalConfiguration.sync({ alter: true });

front/components/assistant/AssistantDetails.tsx

@@ -6,23 +6,29 @@ import {
   CommandLineIcon,
   Modal,
   PlusIcon,
+  ServerIcon,
   TrashIcon,
   XMarkIcon,
 } from "@dust-tt/sparkle";
 import type {
+  AgentConfigurationType,
   AgentUsageType,
   AgentUserListStatus,
   ConnectorProvider,
+  CoreAPITable,
+  DataSourceConfiguration,
+  DustAppRunConfigurationType,
   LightAgentConfigurationType,
+  TablesQueryConfigurationType,
+  WorkspaceType,
+} from "@dust-tt/types";
+import {
+  isDustAppRunConfiguration,
+  isRetrievalConfiguration,
+  isTablesQueryConfiguration,
 } from "@dust-tt/types";
-import type { DustAppRunConfigurationType } from "@dust-tt/types";
-import type { DataSourceConfiguration } from "@dust-tt/types";
-import type { AgentConfigurationType } from "@dust-tt/types";
-import type { WorkspaceType } from "@dust-tt/types";
-import { isDustAppRunConfiguration } from "@dust-tt/types";
-import { isRetrievalConfiguration } from "@dust-tt/types";
 import Link from "next/link";
-import { useContext, useState } from "react";
+import { useCallback, useContext, useEffect, useState } from "react";
 import ReactMarkdown from "react-markdown";
 
 import { DeleteAssistantDialog } from "@app/components/assistant/AssistantActions";
@@ -125,6 +131,11 @@ export function AssistantDetails({
           </div>
           <DataSourcesSection dataSourceConfigurations={action.dataSources} />
         </div>
+      ) : isTablesQueryConfiguration(action) ? (
+        <div className="flex flex-col gap-2">
+          <div className="text-lg font-bold text-element-800">Tables</div>
+          <TablesQuerySection tablesQueryConfig={action} />
+        </div>
       ) : null
     ) : null;
 
@@ -379,3 +390,82 @@ function ButtonsSection({
     </Button.List>
   );
 }
+
+function TablesQuerySection({
+  tablesQueryConfig,
+}: {
+  tablesQueryConfig: TablesQueryConfigurationType;
+}) {
+  const [tables, setTables] = useState<CoreAPITable[] | null>(null);
+  const [isLoading, setIsLoading] = useState<boolean>(false);
+  const [isError, setIsError] = useState<boolean>(false);
+
+  const getTables = useCallback(async () => {
+    const tableEndpoints = tablesQueryConfig.tables.map(
+      (t) =>
+        `/api/w/${t.workspaceId}/data_sources/${t.dataSourceId}/tables/${t.tableId}`
+    );
+
+    const results = await Promise.all(
+      tableEndpoints.map((endpoint) =>
+        fetch(endpoint, {
+          method: "GET",
+          headers: { "Content-Type": "application/json" },
+        })
+      )
+    );
+
+    const tablesParsed = [];
+    for (const res of results) {
+      if (!res.ok) {
+        throw new Error((await res.json()).error.message);
+      }
+      tablesParsed.push((await res.json()).table);
+    }
+
+    setTables(tablesParsed);
+  }, [tablesQueryConfig.tables]);
+
+  useEffect(() => {
+    if (!tablesQueryConfig.tables || isLoading || isError || tables?.length) {
+      return;
+    }
+    setIsLoading(true);
+    getTables()
+      .catch(() => setIsError(true))
+      .finally(() => setIsLoading(false));
+  }, [getTables, isLoading, tablesQueryConfig.tables, tables?.length, isError]);
+
+  if (isLoading) {
+    return (
+      <div className="flex flex-col gap-2">
+        <div>Loading...</div>;
+      </div>
+    );
+  }
+
+  if (!tables) {
+    return (
+      <div className="flex flex-col gap-2">
+        <div>Error loading tables.</div>;
+      </div>
+    );
+  }
+
+  return (
+    <div className="flex flex-col gap-2">
+      <div>The following tables are queried before answering:</div>
+      {tables.map((t) => (
+        <div
+          className="flex flex-row items-center gap-2"
+          key={`${t.data_source_id}/${t.table_id}`}
+        >
+          <div>
+            <ServerIcon />
+          </div>
+          <div key={`${t.data_source_id}/${t.table_id}`}>{t.name}</div>
+        </div>
+      ))}
+    </div>
+  );
+}

front/components/assistant/conversation/AgentAction.tsx

@@ -1,13 +1,13 @@
 import type { AgentActionType } from "@dust-tt/types";
 import {
-  isDatabaseQueryActionType,
   isDustAppRunActionType,
+  isRetrievalActionType,
+  isTablesQueryActionType,
 } from "@dust-tt/types";
-import { isRetrievalActionType } from "@dust-tt/types";
 
-import DatabaseQueryAction from "@app/components/assistant/conversation/DatabaseQueryAction";
 import DustAppRunAction from "@app/components/assistant/conversation/DustAppRunAction";
 import RetrievalAction from "@app/components/assistant/conversation/RetrievalAction";
+import TablesQueryAction from "@app/components/assistant/conversation/TablesQueryAction";
 
 export function AgentAction({ action }: { action: AgentActionType }) {
   if (isRetrievalActionType(action)) {
@@ -24,10 +24,10 @@ export function AgentAction({ action }: { action: AgentActionType }) {
       </div>
     );
   }
-  if (isDatabaseQueryActionType(action)) {
+  if (isTablesQueryActionType(action)) {
     return (
       <div className="pb-4">
-        <DatabaseQueryAction databaseQueryAction={action} />
+        <TablesQueryAction tablesQueryAction={action} />
       </div>
     );
   }

front/components/assistant/conversation/AgentMessage.tsx

@@ -128,8 +128,8 @@ export function AgentMessage({
       case "retrieval_params":
       case "dust_app_run_params":
       case "dust_app_run_block":
-      case "database_query_params":
-      case "database_query_output":
+      case "tables_query_params":
+      case "tables_query_output":
         setStreamedAgentMessage((m) => {
           return { ...m, action: event.action };
         });

front/components/assistant/conversation/DatabaseQueryAction.tsx → front/components/assistant/conversation/TablesQueryAction.tsx

@@ -6,7 +6,7 @@ import {
   Spinner,
   Tooltip,
 } from "@dust-tt/sparkle";
-import type { DatabaseQueryActionType } from "@dust-tt/types";
+import type { TablesQueryActionType } from "@dust-tt/types";
 import dynamic from "next/dynamic";
 import { useState } from "react";
 import { amber, emerald, slate } from "tailwindcss/colors";
@@ -16,20 +16,20 @@ const SyntaxHighlighter = dynamic(
   { ssr: false }
 );
 
-export default function DatabaseQueryAction({
-  databaseQueryAction,
+export default function TablesQueryAction({
+  tablesQueryAction,
 }: {
-  databaseQueryAction: DatabaseQueryActionType;
+  tablesQueryAction: TablesQueryActionType;
 }) {
   const [isOutputExpanded, setIsOutputExpanded] = useState(false);
 
   // Extracting question from the params
-  const params = databaseQueryAction.params;
+  const params = tablesQueryAction.params;
   const question =
     typeof params?.question === "string" ? params.question : null;
 
   // Extracting query and result from the output
-  const output = databaseQueryAction.output;
+  const output = tablesQueryAction.output;
   const query = typeof output?.query === "string" ? output.query : null;
   const noQuery = output?.no_query === true;
   const results = output?.results;