Assistant Builder: Tree component for datasources (#4065)

* Assistant Builder: Tree component for datasources * Use new titleWithParentsContext as default title if set * Add urls to documents in tree * Fix isSelectAll logic
dust-tt · Mar 5, 2024 · 8038623 · 8038623
1 parent d07a04b
commit 8038623
Show file tree

Hide file tree

Showing 10 changed files with 314 additions and 213 deletions.
diff --git a/front/components/DataSourceResourceSelectorTree.tsx b/front/components/DataSourceResourceSelectorTree.tsx
@@ -6,7 +6,11 @@ import {
   DocumentTextIcon,
   Spinner,
 } from "@dust-tt/sparkle";
-import type { DataSourceType, WorkspaceType } from "@dust-tt/types";
+import type {
+  ConnectorNode,
+  DataSourceType,
+  WorkspaceType,
+} from "@dust-tt/types";
 import type { ConnectorNodeType, ConnectorPermission } from "@dust-tt/types";
 import { CircleStackIcon, FolderIcon } from "@heroicons/react/20/solid";
 import { useState } from "react";
@@ -29,7 +33,8 @@ export default function DataSourceResourceSelectorTree({
   expandable: boolean;
   selectedParentIds: Set<string>;
   onSelectChange: (
-    resource: { resourceId: string; resourceName: string; parents: string[] },
+    resource: ConnectorNode,
+    parents: string[],
     selected: boolean
   ) => void;
   parentsById: Record<string, Set<string>>;
@@ -99,7 +104,8 @@ function DataSourceResourceSelectorChildren({
   selectedParentIds: Set<string>;
   parents: string[];
   onSelectChange: (
-    resource: { resourceId: string; resourceName: string; parents: string[] },
+    resource: ConnectorNode,
+    parents: string[],
     selected: boolean
   ) => void;
   parentsById: Record<string, Set<string>>;
@@ -201,14 +207,7 @@ function DataSourceResourceSelectorChildren({
                       checked={checkStatus === "checked"}
                       partialChecked={checkStatus === "partial"}
                       onChange={(checked) =>
-                        onSelectChange(
-                          {
-                            resourceId: r.internalId,
-                            resourceName: r.title,
-                            parents: parents,
-                          },
-                          checked
-                        )
+                        onSelectChange(r, parents, checked)
                       }
                       disabled={isChecked || fullySelected}
                     />

diff --git a/front/components/ManagedDataSourceDocumentModal.tsx b/front/components/ManagedDataSourceDocumentModal.tsx
@@ -10,7 +10,7 @@ export default function ManagedDataSourceDocumentModal({
   setOpen,
 }: {
   owner: WorkspaceType;
-  dataSource: DataSourceType;
+  dataSource: DataSourceType | null;
   documentId: string | null;
   isOpen: boolean;
   setOpen: (open: boolean) => void;
@@ -20,7 +20,7 @@ export default function ManagedDataSourceDocumentModal({
   const [downloading, setDownloading] = useState(false);
 
   useEffect(() => {
-    if (documentId) {
+    if (documentId && dataSource?.name) {
       setDownloading(true);
       fetch(
         `/api/w/${owner.sId}/data_sources/${encodeURIComponent(
@@ -42,7 +42,7 @@ export default function ManagedDataSourceDocumentModal({
         })
         .catch((e) => console.error(e));
     }
-  }, [dataSource.name, documentId, owner.sId]);
+  }, [dataSource, documentId, owner.sId]);
 
   function closeModal() {
     setOpen(false);

diff --git a/front/components/assistant/SlackIntegration.tsx b/front/components/assistant/SlackIntegration.tsx
@@ -71,7 +71,7 @@ export function SlackIntegration({
             dataSource={slackDataSource}
             selectedParentIds={selectedChannelIds}
             parentsById={{}}
-            onSelectChange={({ resourceId, resourceName }, selected) => {
+            onSelectChange={(node, selected) => {
               setHasChanged(true);
 
               if (selected) {
@@ -81,8 +81,8 @@ export function SlackIntegration({
                     finalState.push(...sel);
                   }
                   finalState.push({
-                    slackChannelId: resourceId,
-                    slackChannelName: resourceName,
+                    slackChannelId: node.internalId,
+                    slackChannelName: node.title,
                   });
 
                   return finalState;
@@ -95,7 +95,7 @@ export function SlackIntegration({
                   }
                   finalState.splice(
                     finalState.findIndex(
-                      (c) => c.slackChannelId === resourceId
+                      (c) => c.slackChannelId === node.internalId
                     ),
                     1
                   );

diff --git a/front/components/assistant_builder/ActionScreen.tsx b/front/components/assistant_builder/ActionScreen.tsx
@@ -30,7 +30,6 @@ import { TIME_FRAME_UNIT_TO_LABEL } from "@app/components/assistant_builder/shar
 import TablesSelectionSection from "@app/components/assistant_builder/TablesSelectionSection";
 import type {
   ActionMode,
-  AssistantBuilderDataSourceConfiguration,
   AssistantBuilderState,
 } from "@app/components/assistant_builder/types";
 import { tableKey } from "@app/lib/client/tables_query";
@@ -142,15 +141,9 @@ export default function ActionScreen({
   timeFrameError: string | null;
 }) {
   const [showDataSourcesModal, setShowDataSourcesModal] = useState(false);
-  const [dataSourceToManage, setDataSourceToManage] =
-    useState<AssistantBuilderDataSourceConfiguration | null>(null);
   const [showDustAppsModal, setShowDustAppsModal] = useState(false);
   const [showTableModal, setShowTableModal] = useState(false);
 
-  const configurableDataSources = dataSources.filter(
-    (dataSource) => !builderState.dataSourceConfigurations[dataSource.name]
-  );
-
   const deleteDataSource = (name: string) => {
     setEdited(true);
     setBuilderState(({ dataSourceConfigurations, ...rest }) => {
@@ -197,7 +190,7 @@ export default function ActionScreen({
   };
 
   const noDataSources =
-    configurableDataSources.length === 0 &&
+    dataSources.length === 0 &&
     Object.keys(builderState.dataSourceConfigurations).length === 0;
   const noDustApp = dustApps.length === 0;
 
@@ -207,9 +200,6 @@ export default function ActionScreen({
         isOpen={showDustAppsModal}
         setOpen={(isOpen) => {
           setShowDustAppsModal(isOpen);
-          if (!isOpen) {
-            setDataSourceToManage(null);
-          }
         }}
         dustApps={dustApps}
         onSave={({ app }) => {
@@ -243,12 +233,9 @@ export default function ActionScreen({
         isOpen={showDataSourcesModal}
         setOpen={(isOpen) => {
           setShowDataSourcesModal(isOpen);
-          if (!isOpen) {
-            setDataSourceToManage(null);
-          }
         }}
         owner={owner}
-        dataSources={configurableDataSources}
+        dataSources={dataSources}
         onSave={({ dataSource, selectedResources, isSelectAll }) => {
           setEdited(true);
           setBuilderState((state) => ({
@@ -263,7 +250,8 @@ export default function ActionScreen({
             },
           }));
         }}
-        dataSourceToManage={dataSourceToManage}
+        onDelete={deleteDataSource}
+        dataSourceConfigurations={builderState.dataSourceConfigurations}
       />
       <div className="flex flex-col gap-4 text-sm text-element-700">
         <div className="flex flex-col gap-2">
@@ -414,7 +402,6 @@ export default function ActionScreen({
                   <DropdownMenu.Items origin="topLeft" width={260}>
                     {SEARCH_MODES.filter((key) => {
                       const flag = SEARCH_MODE_SPECIFICATIONS[key].flag;
-                      console.log(owner.flags);
                       return flag === null || owner.flags.includes(flag);
                     }).map((key) => (
                       <DropdownMenu.Item
@@ -451,17 +438,12 @@ export default function ActionScreen({
           }
         >
           <DataSourceSelectionSection
+            owner={owner}
             dataSourceConfigurations={builderState.dataSourceConfigurations}
             openDataSourceModal={() => {
               setShowDataSourcesModal(true);
             }}
-            canAddDataSource={configurableDataSources.length > 0}
-            onManageDataSource={(name) => {
-              setDataSourceToManage(
-                builderState.dataSourceConfigurations[name]
-              );
-              setShowDataSourcesModal(true);
-            }}
+            canAddDataSource={dataSources.length > 0}
             onDelete={deleteDataSource}
           />
         </ActionModeSection>
@@ -472,17 +454,12 @@ export default function ActionScreen({
           }
         >
           <DataSourceSelectionSection
+            owner={owner}
             dataSourceConfigurations={builderState.dataSourceConfigurations}
             openDataSourceModal={() => {
               setShowDataSourcesModal(true);
             }}
-            canAddDataSource={configurableDataSources.length > 0}
-            onManageDataSource={(name) => {
-              setDataSourceToManage(
-                builderState.dataSourceConfigurations[name]
-              );
-              setShowDataSourcesModal(true);
-            }}
+            canAddDataSource={dataSources.length > 0}
             onDelete={deleteDataSource}
           />
           <div className={"flex flex-row items-center gap-4 pb-4"}>

diff --git a/front/components/assistant_builder/AssistantBuilder.tsx b/front/components/assistant_builder/AssistantBuilder.tsx
@@ -603,7 +603,12 @@ async function submitForm({
             workspaceId: owner.sId,
             filter: {
               parents: !isSelectAll
-                ? { in: Object.keys(selectedResources), not: [] }
+                ? {
+                    in: selectedResources.map(
+                      (resource) => resource.internalId
+                    ),
+                    not: [],
+                  }
                 : null,
               tags: null,
             },