From 9f60d6c47458947f8dd0ffe1f41d7639ebe97914 Mon Sep 17 00:00:00 2001
From: Henry Fontanier <h.fontanier@gmail.com>
Date: Thu, 18 Jan 2024 12:05:48 +0100
Subject: [PATCH] enh: use table names in assistant builder (#3286)

* use full tables in builder state

* display all datasources even if there is a dsconfig

* move files into the /components

* dont refetch DS

* lint

---------

Co-authored-by: Henry Fontanier <henry@dust.tt>
---
 .../assistant_builder/AssistantBuilder.tsx    |  94 ++---------
 .../AssistantBuilderDataSourceModal.tsx       |   2 +-
 .../AssistantBuilderDustAppModal.tsx          |   2 +-
 .../AssistantBuilderTablesModal.tsx           |   3 +-
 .../DataSourceSelectionSection.tsx            |   2 +-
 .../DustAppSelectionSection.tsx               |   2 +-
 .../TablesSelectionSection.tsx                |  11 +-
 .../server_side_props_helpers.ts              | 151 ++++++++++++++++++
 front/components/assistant_builder/types.ts   |  78 +++++++++
 front/lib/client/tables_query.ts              |   7 +
 .../[wId]/builder/assistants/[aId]/index.tsx  | 127 ++-------------
 .../pages/w/[wId]/builder/assistants/new.tsx  |  20 +--
 12 files changed, 279 insertions(+), 220 deletions(-)
 create mode 100644 front/components/assistant_builder/server_side_props_helpers.ts
 create mode 100644 front/components/assistant_builder/types.ts
 create mode 100644 front/lib/client/tables_query.ts

diff --git a/front/components/assistant_builder/AssistantBuilder.tsx b/front/components/assistant_builder/AssistantBuilder.tsx
index b6483283cd3f..0134de2ec5e8 100644
--- a/front/components/assistant_builder/AssistantBuilder.tsx
+++ b/front/components/assistant_builder/AssistantBuilder.tsx
@@ -59,6 +59,14 @@ import {
 } from "@app/components/assistant_builder/shared";
 import TablesSelectionSection from "@app/components/assistant_builder/TablesSelectionSection";
 import { TeamSharingSection } from "@app/components/assistant_builder/TeamSharingSection";
+import {
+  type ActionMode,
+  ADVANCED_ACTION_MODES,
+  type AssistantBuilderDataSourceConfiguration,
+  type AssistantBuilderInitialState,
+  type AssistantBuilderState,
+  BASIC_ACTION_MODES,
+} from "@app/components/assistant_builder/types";
 import DataSourceResourceSelectorTree from "@app/components/DataSourceResourceSelectorTree";
 import AppLayout from "@app/components/sparkle/AppLayout";
 import {
@@ -68,6 +76,7 @@ import {
 import { subNavigationAssistants } from "@app/components/sparkle/navigation";
 import { SendNotificationsContext } from "@app/components/sparkle/Notification";
 import { getSupportedModelConfig } from "@app/lib/assistant";
+import { tableKey } from "@app/lib/client/tables_query";
 import { CONNECTOR_CONFIGURATIONS } from "@app/lib/connector_providers";
 import { isActivatedStructuredDB } from "@app/lib/development";
 import { isUpgraded } from "@app/lib/plans/plan_codes";
@@ -98,19 +107,6 @@ const SPIRIT_AVATAR_URLS = SPIRIT_AVATAR_FILES.map((f) =>
   buildAvatarUrl(SPIRIT_AVATARS_BASE_PATH, f)
 );
 
-// Actions
-
-const BASIC_ACTION_MODES = ["GENERIC", "RETRIEVAL_SEARCH"] as const;
-const ADVANCED_ACTION_MODES = [
-  "RETRIEVAL_EXHAUSTIVE",
-  "DUST_APP_RUN",
-  "TABLES_QUERY",
-] as const;
-
-type ActionMode =
-  | (typeof BASIC_ACTION_MODES)[number]
-  | (typeof ADVANCED_ACTION_MODES)[number];
-
 const ACTION_MODE_TO_LABEL: Record<ActionMode, string> = {
   GENERIC: "No action",
   RETRIEVAL_SEARCH: "Search in data sources",
@@ -137,74 +133,6 @@ export const CONNECTOR_PROVIDER_TO_RESOURCE_NAME: Record<
   webcrawler: { singular: "page", plural: "pages" },
 };
 
-export type AssistantBuilderDataSourceConfiguration = {
-  dataSource: DataSourceType;
-  selectedResources: Record<string, string>;
-  isSelectAll: boolean;
-};
-
-// DustAppRun Action
-
-export type AssistantBuilderDustAppConfiguration = {
-  app: AppType;
-};
-
-// Tables Query Action
-
-export type AssistantBuilderTableConfiguration = {
-  dataSourceId: string;
-  workspaceId: string;
-  tableId: string;
-};
-
-// Builder State
-
-type AssistantBuilderState = {
-  actionMode: ActionMode;
-  dataSourceConfigurations: Record<
-    string,
-    AssistantBuilderDataSourceConfiguration
-  >;
-  timeFrame: {
-    value: number;
-    unit: TimeframeUnit;
-  };
-  dustAppConfiguration: AssistantBuilderDustAppConfiguration | null;
-  tablesQueryConfiguration: Record<string, AssistantBuilderTableConfiguration>;
-  handle: string | null;
-  description: string | null;
-  scope: Exclude<AgentConfigurationScope, "global">;
-  instructions: string | null;
-  avatarUrl: string | null;
-  generationSettings: {
-    modelSettings: SupportedModel;
-    temperature: number;
-  };
-};
-
-// initial state is like the state, but:
-// - doesn't allow null handle/description/instructions
-// - allows null timeFrame
-// - allows null dataSourceConfigurations
-export type AssistantBuilderInitialState = {
-  actionMode: AssistantBuilderState["actionMode"];
-  dataSourceConfigurations:
-    | AssistantBuilderState["dataSourceConfigurations"]
-    | null;
-  timeFrame: AssistantBuilderState["timeFrame"] | null;
-  dustAppConfiguration: AssistantBuilderState["dustAppConfiguration"];
-  tablesQueryConfiguration: AssistantBuilderState["tablesQueryConfiguration"];
-  handle: string;
-  description: string;
-  scope: Exclude<AgentConfigurationScope, "global">;
-  instructions: string;
-  avatarUrl: string | null;
-  generationSettings: {
-    modelSettings: SupportedModel;
-    temperature: number;
-  } | null;
-};
-
 export const BUILDER_FLOWS = [
   "workspace_assistants",
   "personal_assistants",
@@ -732,14 +660,14 @@ export default function AssistantBuilder({
         isOpen={showTableModal}
         setOpen={(isOpen) => setShowTableModal(isOpen)}
         owner={owner}
-        dataSources={configurableDataSources}
+        dataSources={dataSources}
         onSave={(t) => {
           setEdited(true);
           setBuilderState((state) => ({
             ...state,
             tablesQueryConfiguration: {
               ...state.tablesQueryConfiguration,
-              [`${t.workspaceId}/${t.dataSourceId}/${t.tableId}`]: t,
+              [tableKey(t)]: t,
             },
           }));
         }}
diff --git a/front/components/assistant_builder/AssistantBuilderDataSourceModal.tsx b/front/components/assistant_builder/AssistantBuilderDataSourceModal.tsx
index 4e41644fa2ee..8f3d594c78a0 100644
--- a/front/components/assistant_builder/AssistantBuilderDataSourceModal.tsx
+++ b/front/components/assistant_builder/AssistantBuilderDataSourceModal.tsx
@@ -13,8 +13,8 @@ import { Transition } from "@headlessui/react";
 import * as React from "react";
 import { useCallback, useEffect, useState } from "react";
 
-import type { AssistantBuilderDataSourceConfiguration } from "@app/components/assistant_builder/AssistantBuilder";
 import { CONNECTOR_PROVIDER_TO_RESOURCE_NAME } from "@app/components/assistant_builder/AssistantBuilder";
+import type { AssistantBuilderDataSourceConfiguration } from "@app/components/assistant_builder/types";
 import DataSourceResourceSelectorTree from "@app/components/DataSourceResourceSelectorTree";
 import { CONNECTOR_CONFIGURATIONS } from "@app/lib/connector_providers";
 import type { GetConnectorResourceParentsResponseBody } from "@app/pages/api/w/[wId]/data_sources/[name]/managed/parents";
diff --git a/front/components/assistant_builder/AssistantBuilderDustAppModal.tsx b/front/components/assistant_builder/AssistantBuilderDustAppModal.tsx
index dc5fe3ec66b8..6c6eeafc3bb1 100644
--- a/front/components/assistant_builder/AssistantBuilderDustAppModal.tsx
+++ b/front/components/assistant_builder/AssistantBuilderDustAppModal.tsx
@@ -8,7 +8,7 @@ import {
 import type { AppType } from "@dust-tt/types";
 import { Transition } from "@headlessui/react";
 
-import type { AssistantBuilderDustAppConfiguration } from "@app/components/assistant_builder/AssistantBuilder";
+import type { AssistantBuilderDustAppConfiguration } from "@app/components/assistant_builder/types";
 
 export default function AssistantBuilderDustAppModal({
   isOpen,
diff --git a/front/components/assistant_builder/AssistantBuilderTablesModal.tsx b/front/components/assistant_builder/AssistantBuilderTablesModal.tsx
index 45e4d71d779e..2c578ca5c032 100644
--- a/front/components/assistant_builder/AssistantBuilderTablesModal.tsx
+++ b/front/components/assistant_builder/AssistantBuilderTablesModal.tsx
@@ -12,7 +12,7 @@ import { Transition } from "@headlessui/react";
 import * as React from "react";
 import { useState } from "react";
 
-import type { AssistantBuilderTableConfiguration } from "@app/components/assistant_builder/AssistantBuilder";
+import type { AssistantBuilderTableConfiguration } from "@app/components/assistant_builder/types";
 import { CONNECTOR_CONFIGURATIONS } from "@app/lib/connector_providers";
 import { useTables } from "@app/lib/swr";
 
@@ -75,6 +75,7 @@ export default function AssistantBuilderTablesModal({
                 workspaceId: owner.sId,
                 dataSourceId: table.data_source_id,
                 tableId: table.table_id,
+                tableName: table.name,
               };
               setSelectedTable(config);
               onSave(config);
diff --git a/front/components/assistant_builder/DataSourceSelectionSection.tsx b/front/components/assistant_builder/DataSourceSelectionSection.tsx
index 19c3d11d1c99..df564dcf06c4 100644
--- a/front/components/assistant_builder/DataSourceSelectionSection.tsx
+++ b/front/components/assistant_builder/DataSourceSelectionSection.tsx
@@ -7,8 +7,8 @@ import {
   TrashIcon,
 } from "@dust-tt/sparkle";
 
-import type { AssistantBuilderDataSourceConfiguration } from "@app/components/assistant_builder/AssistantBuilder";
 import { CONNECTOR_PROVIDER_TO_RESOURCE_NAME } from "@app/components/assistant_builder/shared";
+import type { AssistantBuilderDataSourceConfiguration } from "@app/components/assistant_builder/types";
 import { EmptyCallToAction } from "@app/components/EmptyCallToAction";
 import { CONNECTOR_CONFIGURATIONS } from "@app/lib/connector_providers";
 import { getDisplayNameForDataSource } from "@app/lib/data_sources";
diff --git a/front/components/assistant_builder/DustAppSelectionSection.tsx b/front/components/assistant_builder/DustAppSelectionSection.tsx
index bb063009d048..e60dd6dd1e43 100644
--- a/front/components/assistant_builder/DustAppSelectionSection.tsx
+++ b/front/components/assistant_builder/DustAppSelectionSection.tsx
@@ -6,7 +6,7 @@ import {
 } from "@dust-tt/sparkle";
 import { Transition } from "@headlessui/react";
 
-import type { AssistantBuilderDustAppConfiguration } from "@app/components/assistant_builder/AssistantBuilder";
+import type { AssistantBuilderDustAppConfiguration } from "@app/components/assistant_builder/types";
 import { EmptyCallToAction } from "@app/components/EmptyCallToAction";
 
 export default function DustAppSelectionSection({
diff --git a/front/components/assistant_builder/TablesSelectionSection.tsx b/front/components/assistant_builder/TablesSelectionSection.tsx
index e66207d0f2dd..1e43b82e803f 100644
--- a/front/components/assistant_builder/TablesSelectionSection.tsx
+++ b/front/components/assistant_builder/TablesSelectionSection.tsx
@@ -7,8 +7,9 @@ import {
 } from "@dust-tt/sparkle";
 import { Transition } from "@headlessui/react";
 
-import type { AssistantBuilderTableConfiguration } from "@app/components/assistant_builder/AssistantBuilder";
+import type { AssistantBuilderTableConfiguration } from "@app/components/assistant_builder/types";
 import { EmptyCallToAction } from "@app/components/EmptyCallToAction";
+import { tableKey } from "@app/lib/client/tables_query";
 
 export default function TablesSelectionSection({
   show,
@@ -69,12 +70,12 @@ export default function TablesSelectionSection({
         ) : (
           <ContextItem.List className="mt-6 border-b border-t border-structure-200">
             {Object.values(tablesQueryConfiguration).map((t) => {
-              const tableKey = `${t.workspaceId}/${t.dataSourceId}/${t.tableId}`;
+              const key = tableKey(t);
               return (
                 <ContextItem
-                  title={t.tableId} // TODO: fetch table name
+                  title={`${t.tableName} (${t.dataSourceId})`}
                   visual={<ContextItem.Visual visual={ServerIcon} />}
-                  key={tableKey}
+                  key={key}
                   action={
                     <Button.List>
                       <Button
@@ -83,7 +84,7 @@ export default function TablesSelectionSection({
                         label="Remove"
                         labelVisible={false}
                         onClick={() => {
-                          onDelete?.(tableKey);
+                          onDelete?.(key);
                         }}
                       />
                     </Button.List>
diff --git a/front/components/assistant_builder/server_side_props_helpers.ts b/front/components/assistant_builder/server_side_props_helpers.ts
new file mode 100644
index 000000000000..9d8543d1c562
--- /dev/null
+++ b/front/components/assistant_builder/server_side_props_helpers.ts
@@ -0,0 +1,151 @@
+import type {
+  AgentConfigurationType,
+  AppType,
+  CoreAPITable,
+  DataSourceType,
+} from "@dust-tt/types";
+import {
+  ConnectorsAPI,
+  CoreAPI,
+  isDustAppRunConfiguration,
+  isRetrievalConfiguration,
+  isTablesQueryConfiguration,
+} from "@dust-tt/types";
+
+import type {
+  AssistantBuilderDataSourceConfiguration,
+  AssistantBuilderInitialState,
+} from "@app/components/assistant_builder/types";
+import { tableKey } from "@app/lib/client/tables_query";
+import logger from "@app/logger/logger";
+
+export async function buildInitialState({
+  dataSourceByName,
+  config,
+  dustApps,
+}: {
+  dataSourceByName: Record<string, DataSourceType>;
+  config: AgentConfigurationType;
+  dustApps: AppType[];
+}) {
+  const coreAPI = new CoreAPI(logger);
+
+  const selectedResources: {
+    dataSourceName: string;
+    resources: string[] | null;
+    isSelectAll: boolean;
+  }[] = [];
+
+  if (isRetrievalConfiguration(config.action)) {
+    for (const ds of config.action.dataSources) {
+      selectedResources.push({
+        dataSourceName: ds.dataSourceId,
+        resources: ds.filter.parents?.in ?? null,
+        isSelectAll: !ds.filter.parents,
+      });
+    }
+  }
+
+  const dataSourceConfigurationsArray: NonNullable<
+    AssistantBuilderInitialState["dataSourceConfigurations"]
+  >[string][] = await Promise.all(
+    selectedResources.map(
+      async (ds): Promise<AssistantBuilderDataSourceConfiguration> => {
+        const dataSource = dataSourceByName[ds.dataSourceName];
+        if (!dataSource.connectorId || !ds.resources) {
+          return {
+            dataSource: dataSource,
+            selectedResources: {},
+            isSelectAll: ds.isSelectAll,
+          };
+        }
+        const connectorsAPI = new ConnectorsAPI(logger);
+        const response = await connectorsAPI.getResourcesTitles({
+          connectorId: dataSource.connectorId,
+          resourceInternalIds: ds.resources,
+        });
+
+        if (response.isErr()) {
+          throw response.error;
+        }
+
+        // key: interalId, value: title
+        const selectedResources: Record<string, string> = {};
+        for (const resource of response.value.resources) {
+          selectedResources[resource.internalId] = resource.title;
+        }
+
+        return {
+          dataSource: dataSource,
+          selectedResources,
+          isSelectAll: ds.isSelectAll,
+        };
+      }
+    )
+  );
+
+  // key: dataSourceName, value: DataSourceConfig
+  const dataSourceConfigurations = dataSourceConfigurationsArray.reduce(
+    (acc, curr) => ({ ...acc, [curr.dataSource.name]: curr }),
+    {} as Record<string, AssistantBuilderDataSourceConfiguration>
+  );
+
+  let dustAppConfiguration: AssistantBuilderInitialState["dustAppConfiguration"] =
+    null;
+
+  if (isDustAppRunConfiguration(config.action)) {
+    for (const app of dustApps) {
+      if (app.sId === config.action.appId) {
+        dustAppConfiguration = {
+          app,
+        };
+        break;
+      }
+    }
+  }
+
+  let tablesQueryConfiguration: AssistantBuilderInitialState["tablesQueryConfiguration"] =
+    {};
+
+  if (
+    isTablesQueryConfiguration(config.action) &&
+    config.action.tables.length
+  ) {
+    const coreAPITables: CoreAPITable[] = await Promise.all(
+      config.action.tables.map(async (t) => {
+        const dataSource = dataSourceByName[t.dataSourceId];
+        const coreAPITable = await coreAPI.getTable({
+          projectId: dataSource.dustAPIProjectId,
+          dataSourceName: dataSource.name,
+          tableId: t.tableId,
+        });
+
+        if (coreAPITable.isErr()) {
+          throw coreAPITable.error;
+        }
+
+        return coreAPITable.value.table;
+      })
+    );
+
+    tablesQueryConfiguration = config.action.tables.reduce((acc, curr, i) => {
+      const table = coreAPITables[i];
+      const key = tableKey(curr);
+      return {
+        ...acc,
+        [key]: {
+          workspaceId: curr.workspaceId,
+          dataSourceId: curr.dataSourceId,
+          tableId: curr.tableId,
+          tableName: `${table.name}`,
+        },
+      };
+    }, {} as AssistantBuilderInitialState["tablesQueryConfiguration"]);
+  }
+
+  return {
+    dataSourceConfigurations,
+    dustAppConfiguration,
+    tablesQueryConfiguration,
+  };
+}
diff --git a/front/components/assistant_builder/types.ts b/front/components/assistant_builder/types.ts
new file mode 100644
index 000000000000..8a459cdc6195
--- /dev/null
+++ b/front/components/assistant_builder/types.ts
@@ -0,0 +1,78 @@
+import type {
+  AgentConfigurationScope,
+  AppType,
+  DataSourceType,
+  SupportedModel,
+  TimeframeUnit,
+} from "@dust-tt/types";
+
+export const BASIC_ACTION_MODES = ["GENERIC", "RETRIEVAL_SEARCH"] as const;
+export const ADVANCED_ACTION_MODES = [
+  "RETRIEVAL_EXHAUSTIVE",
+  "DUST_APP_RUN",
+  "TABLES_QUERY",
+] as const;
+
+export type ActionMode =
+  | (typeof BASIC_ACTION_MODES)[number]
+  | (typeof ADVANCED_ACTION_MODES)[number];
+
+export type AssistantBuilderDataSourceConfiguration = {
+  dataSource: DataSourceType;
+  selectedResources: Record<string, string>;
+  isSelectAll: boolean;
+};
+
+export type AssistantBuilderDustAppConfiguration = {
+  app: AppType;
+};
+
+export type AssistantBuilderTableConfiguration = {
+  dataSourceId: string;
+  workspaceId: string;
+  tableId: string;
+  tableName: string;
+};
+
+// Builder State
+export type AssistantBuilderState = {
+  actionMode: ActionMode;
+  dataSourceConfigurations: Record<
+    string,
+    AssistantBuilderDataSourceConfiguration
+  >;
+  timeFrame: {
+    value: number;
+    unit: TimeframeUnit;
+  };
+  dustAppConfiguration: AssistantBuilderDustAppConfiguration | null;
+  tablesQueryConfiguration: Record<string, AssistantBuilderTableConfiguration>;
+  handle: string | null;
+  description: string | null;
+  scope: Exclude<AgentConfigurationScope, "global">;
+  instructions: string | null;
+  avatarUrl: string | null;
+  generationSettings: {
+    modelSettings: SupportedModel;
+    temperature: number;
+  };
+};
+
+export type AssistantBuilderInitialState = {
+  actionMode: AssistantBuilderState["actionMode"];
+  dataSourceConfigurations:
+    | AssistantBuilderState["dataSourceConfigurations"]
+    | null;
+  timeFrame: AssistantBuilderState["timeFrame"] | null;
+  dustAppConfiguration: AssistantBuilderState["dustAppConfiguration"];
+  tablesQueryConfiguration: AssistantBuilderState["tablesQueryConfiguration"];
+  handle: string;
+  description: string;
+  scope: Exclude<AgentConfigurationScope, "global">;
+  instructions: string;
+  avatarUrl: string | null;
+  generationSettings: {
+    modelSettings: SupportedModel;
+    temperature: number;
+  } | null;
+};
diff --git a/front/lib/client/tables_query.ts b/front/lib/client/tables_query.ts
new file mode 100644
index 000000000000..427f313ee123
--- /dev/null
+++ b/front/lib/client/tables_query.ts
@@ -0,0 +1,7 @@
+import type { TablesQueryConfigurationType } from "@dust-tt/types";
+
+export function tableKey(
+  table: TablesQueryConfigurationType["tables"][number]
+) {
+  return `${table.workspaceId}/${table.dataSourceId}/${table.tableId}`;
+}
diff --git a/front/pages/w/[wId]/builder/assistants/[aId]/index.tsx b/front/pages/w/[wId]/builder/assistants/[aId]/index.tsx
index 69e64f56829b..ddbe67a5575a 100644
--- a/front/pages/w/[wId]/builder/assistants/[aId]/index.tsx
+++ b/front/pages/w/[wId]/builder/assistants/[aId]/index.tsx
@@ -8,138 +8,28 @@ import type {
   WorkspaceType,
 } from "@dust-tt/types";
 import {
-  ConnectorsAPI,
   isDustAppRunConfiguration,
   isRetrievalConfiguration,
   isTablesQueryConfiguration,
 } from "@dust-tt/types";
 import type { GetServerSideProps, InferGetServerSidePropsType } from "next";
 
-import type {
-  AssistantBuilderInitialState,
-  BuilderFlow,
-} from "@app/components/assistant_builder/AssistantBuilder";
+import type { BuilderFlow } from "@app/components/assistant_builder/AssistantBuilder";
 import AssistantBuilder, {
   BUILDER_FLOWS,
 } from "@app/components/assistant_builder/AssistantBuilder";
+import { buildInitialState } from "@app/components/assistant_builder/server_side_props_helpers";
+import type {
+  AssistantBuilderDataSourceConfiguration,
+  AssistantBuilderInitialState,
+} from "@app/components/assistant_builder/types";
 import { getApps } from "@app/lib/api/app";
 import { getAgentConfiguration } from "@app/lib/api/assistant/configuration";
 import { getDataSources } from "@app/lib/api/data_sources";
 import { Authenticator, getSession, getUserFromSession } from "@app/lib/auth";
-import logger from "@app/logger/logger";
 
 const { GA_TRACKING_ID = "" } = process.env;
 
-type DataSourceConfig = NonNullable<
-  AssistantBuilderInitialState["dataSourceConfigurations"]
->[string];
-
-export async function buildInitialState({
-  dataSourceByName,
-  config,
-  dustApps,
-}: {
-  dataSourceByName: Record<string, DataSourceType>;
-  config: AgentConfigurationType;
-  dustApps: AppType[];
-}) {
-  const selectedResources: {
-    dataSourceName: string;
-    resources: string[] | null;
-    isSelectAll: boolean;
-  }[] = [];
-
-  if (isRetrievalConfiguration(config.action)) {
-    for (const ds of config.action.dataSources) {
-      selectedResources.push({
-        dataSourceName: ds.dataSourceId,
-        resources: ds.filter.parents?.in ?? null,
-        isSelectAll: !ds.filter.parents,
-      });
-    }
-  }
-
-  const dataSourceConfigurationsArray: DataSourceConfig[] = await Promise.all(
-    selectedResources.map(async (ds): Promise<DataSourceConfig> => {
-      const dataSource = dataSourceByName[ds.dataSourceName];
-      if (!dataSource.connectorId || !ds.resources) {
-        return {
-          dataSource: dataSource,
-          selectedResources: {},
-          isSelectAll: ds.isSelectAll,
-        };
-      }
-      const connectorsAPI = new ConnectorsAPI(logger);
-      const response = await connectorsAPI.getResourcesTitles({
-        connectorId: dataSource.connectorId,
-        resourceInternalIds: ds.resources,
-      });
-
-      if (response.isErr()) {
-        throw response.error;
-      }
-
-      // key: interalId, value: title
-      const selectedResources: Record<string, string> = {};
-      for (const resource of response.value.resources) {
-        selectedResources[resource.internalId] = resource.title;
-      }
-
-      return {
-        dataSource: dataSource,
-        selectedResources,
-        isSelectAll: ds.isSelectAll,
-      };
-    })
-  );
-
-  // key: dataSourceName, value: DataSourceConfig
-  const dataSourceConfigurations = dataSourceConfigurationsArray.reduce(
-    (acc, curr) => ({ ...acc, [curr.dataSource.name]: curr }),
-    {} as Record<string, DataSourceConfig>
-  );
-
-  let dustAppConfiguration: AssistantBuilderInitialState["dustAppConfiguration"] =
-    null;
-
-  if (isDustAppRunConfiguration(config.action)) {
-    for (const app of dustApps) {
-      if (app.sId === config.action.appId) {
-        dustAppConfiguration = {
-          app,
-        };
-        break;
-      }
-    }
-  }
-
-  let tablesQueryConfiguration: AssistantBuilderInitialState["tablesQueryConfiguration"] =
-    {};
-
-  if (
-    isTablesQueryConfiguration(config.action) &&
-    config.action.tables.length
-  ) {
-    tablesQueryConfiguration = config.action.tables.reduce(
-      (acc, curr) => ({
-        ...acc,
-        [`${curr.workspaceId}/${curr.dataSourceId}/${curr.tableId}`]: {
-          workspaceId: curr.workspaceId,
-          dataSourceId: curr.dataSourceId,
-          tableId: curr.tableId,
-        },
-      }),
-      {} as AssistantBuilderInitialState["tablesQueryConfiguration"]
-    );
-  }
-
-  return {
-    dataSourceConfigurations,
-    dustAppConfiguration,
-    tablesQueryConfiguration,
-  };
-}
-
 export const getServerSideProps: GetServerSideProps<{
   user: UserType;
   owner: WorkspaceType;
@@ -147,7 +37,10 @@ export const getServerSideProps: GetServerSideProps<{
   plan: PlanType;
   gaTrackingId: string;
   dataSources: DataSourceType[];
-  dataSourceConfigurations: Record<string, DataSourceConfig>;
+  dataSourceConfigurations: Record<
+    string,
+    AssistantBuilderDataSourceConfiguration
+  >;
   dustApps: AppType[];
   dustAppConfiguration: AssistantBuilderInitialState["dustAppConfiguration"];
   tablesQueryConfiguration: AssistantBuilderInitialState["tablesQueryConfiguration"];
diff --git a/front/pages/w/[wId]/builder/assistants/new.tsx b/front/pages/w/[wId]/builder/assistants/new.tsx
index d49b85887b6f..ccc32d593914 100644
--- a/front/pages/w/[wId]/builder/assistants/new.tsx
+++ b/front/pages/w/[wId]/builder/assistants/new.tsx
@@ -14,25 +14,22 @@ import {
 } from "@dust-tt/types";
 import type { GetServerSideProps, InferGetServerSidePropsType } from "next";
 
-import type {
-  AssistantBuilderInitialState,
-  BuilderFlow,
-} from "@app/components/assistant_builder/AssistantBuilder";
+import type { BuilderFlow } from "@app/components/assistant_builder/AssistantBuilder";
 import AssistantBuilder, {
   BUILDER_FLOWS,
 } from "@app/components/assistant_builder/AssistantBuilder";
+import { buildInitialState } from "@app/components/assistant_builder/server_side_props_helpers";
+import type {
+  AssistantBuilderDataSourceConfiguration,
+  AssistantBuilderInitialState,
+} from "@app/components/assistant_builder/types";
 import { getApps } from "@app/lib/api/app";
 import { getAgentConfiguration } from "@app/lib/api/assistant/configuration";
 import { getDataSources } from "@app/lib/api/data_sources";
 import { Authenticator, getSession, getUserFromSession } from "@app/lib/auth";
-import { buildInitialState } from "@app/pages/w/[wId]/builder/assistants/[aId]";
 
 const { GA_TRACKING_ID = "" } = process.env;
 
-type DataSourceConfig = NonNullable<
-  AssistantBuilderInitialState["dataSourceConfigurations"]
->[string];
-
 export const getServerSideProps: GetServerSideProps<{
   user: UserType;
   owner: WorkspaceType;
@@ -40,7 +37,10 @@ export const getServerSideProps: GetServerSideProps<{
   plan: PlanType;
   gaTrackingId: string;
   dataSources: DataSourceType[];
-  dataSourceConfigurations: Record<string, DataSourceConfig> | null;
+  dataSourceConfigurations: Record<
+    string,
+    AssistantBuilderDataSourceConfiguration
+  > | null;
   dustApps: AppType[];
   dustAppConfiguration: AssistantBuilderInitialState["dustAppConfiguration"];
   tablesQueryConfiguration: AssistantBuilderInitialState["tablesQueryConfiguration"];