Refactor Assistant builder: use Vaults, use Context

dust-tt · Aug 22, 2024 · 581a3d1 · 581a3d1
1 parent 16268a3
commit 581a3d1
Show file tree

Hide file tree

Showing 58 changed files with 1,024 additions and 965 deletions.
diff --git a/.husky/pre-commit b/.husky/pre-commit
@@ -30,6 +30,18 @@ for directory in "${directories[@]}"; do
             echo "Linting failed. Please fix the issues before committing."
             exit 1
         fi
+        if ! npm run tsc --prefix $root_path/$directory; then
+            echo "Type checking failed. Please fix the issues before committing."
+            exit 1
+        fi
+        if ! npm run format:check --prefix $root_path/$directory; then
+            echo "Formatting check failed. Please fix the issues before committing."
+            exit 1
+        fi
+        if ! npm run docs:check --prefix $root_path/$directory; then
+            echo "Documentation check failed. Please fix the issues before committing."
+            exit 1
+        fi
         break
     fi
 done

diff --git a/front/components/ConnectorPermissionsTree.tsx b/front/components/ConnectorPermissionsTree.tsx
@@ -11,7 +11,7 @@ import {
 import type {
   ConnectorProvider,
   DataSourceType,
-  WorkspaceType,
+  LightWorkspaceType,
 } from "@dust-tt/types";
 import type { ConnectorPermission } from "@dust-tt/types";
 import { useState } from "react";
@@ -62,7 +62,7 @@ export function PermissionTreeChildren({
   useConnectorPermissionsHook,
   isSearchEnabled,
 }: {
-  owner: WorkspaceType;
+  owner: LightWorkspaceType;
   dataSource: DataSourceType;
   parentId: string | null;
   permissionFilter?: ConnectorPermission;
@@ -89,7 +89,7 @@ export function PermissionTreeChildren({
   const { resources, isResourcesLoading, isResourcesError } =
     useConnectorPermissionsHook({
       owner,
-      dataSourceOrView: dataSource,
+      dataSource,
       parentId,
       filterPermission: permissionFilter || null,
     });
@@ -290,7 +290,7 @@ export function PermissionTree({
   showExpand,
   isSearchEnabled,
 }: {
-  owner: WorkspaceType;
+  owner: LightWorkspaceType;
   dataSource: DataSourceType;
   permissionFilter?: ConnectorPermission;
   canUpdatePermissions?: boolean;

diff --git a/front/components/DataSourceResourceSelectorTree.tsx b/front/components/DataSourceResourceSelectorTree.tsx
@@ -4,21 +4,20 @@ import {
   Tree,
 } from "@dust-tt/sparkle";
 import type {
-  ContentNode,
   ContentNodesViewType,
-  DataSourceType,
   DataSourceViewType,
-  WorkspaceType,
+  LightContentNode,
+  LightWorkspaceType,
 } from "@dust-tt/types";
 import type { ConnectorPermission, ContentNodeType } from "@dust-tt/types";
 import { CircleStackIcon, FolderIcon } from "@heroicons/react/20/solid";
 import { useEffect } from "react";
 
-import { useConnectorPermissions } from "@app/lib/swr";
+import { useVaultDataSourceViewContent } from "@app/lib/swr";
 
 export default function DataSourceResourceSelectorTree({
   owner,
-  dataSourceOrView,
+  dataSourceView,
   showExpand, //if not, it's flat
   parentIsSelected,
   selectedParents = [],
@@ -27,14 +26,14 @@ export default function DataSourceResourceSelectorTree({
   filterPermission = "read",
   viewType = "documents",
 }: {
-  owner: WorkspaceType;
-  dataSourceOrView: DataSourceType | DataSourceViewType;
+  owner: LightWorkspaceType;
+  dataSourceView: DataSourceViewType;
   showExpand: boolean;
   parentIsSelected?: boolean;
   selectedParents?: string[];
   selectedResourceIds: string[];
   onSelectChange: (
-    resource: ContentNode,
+    resource: LightContentNode,
     parents: string[],
     selected: boolean
   ) => void;
@@ -45,7 +44,7 @@ export default function DataSourceResourceSelectorTree({
     <div className="overflow-x-auto">
       <DataSourceResourceSelectorChildren
         owner={owner}
-        dataSourceOrView={dataSourceOrView}
+        dataSourceView={dataSourceView}
         parentId={null}
         showExpand={showExpand}
         parents={[]}
@@ -87,7 +86,7 @@ function getIconForType(type: ContentNodeType): IconComponentType {
 
 function DataSourceResourceSelectorChildren({
   owner,
-  dataSourceOrView,
+  dataSourceView,
   parentId,
   parents,
   parentIsSelected,
@@ -98,50 +97,51 @@ function DataSourceResourceSelectorChildren({
   filterPermission,
   viewType = "documents",
 }: {
-  owner: WorkspaceType;
-  dataSourceOrView: DataSourceType | DataSourceViewType;
+  owner: LightWorkspaceType;
+  dataSourceView: DataSourceViewType;
   parentId: string | null;
   parents: string[];
   parentIsSelected?: boolean;
   selectedParents: string[];
   showExpand: boolean;
   selectedResourceIds: string[];
   onSelectChange: (
-    resource: ContentNode,
+    resource: LightContentNode,
     parents: string[],
     selected: boolean
   ) => void;
   filterPermission: ConnectorPermission;
   viewType: ContentNodesViewType;
 }) {
-  const { resources, isResourcesLoading, isResourcesError } =
-    useConnectorPermissions({
-      owner: owner,
-      dataSourceOrView,
+  const { vaultContent, isVaultContentLoading, isVaultContentError } =
+    useVaultDataSourceViewContent({
+      workspaceId: owner.sId,
+      vaultId: dataSourceView.vaultId,
+      dataSourceViewId: dataSourceView.sId,
+      viewType,
       parentId,
       filterPermission,
-      disabled: dataSourceOrView.connectorId === null,
-      viewType,
+      disabled: dataSourceView.dataSource.connectorId === null,
     });
 
   useEffect(() => {
     if (parentIsSelected) {
       // Unselected previously selected children
-      resources
+      vaultContent
         .filter((r) => selectedResourceIds.includes(r.internalId))
         .forEach((r) => {
           onSelectChange(r, parents, false);
         });
     }
   }, [
-    resources,
+    vaultContent,
     parentIsSelected,
     selectedResourceIds,
     onSelectChange,
     parents,
   ]);
 
-  if (isResourcesError) {
+  if (isVaultContentError) {
     return (
       <div className="text-warning text-sm">
         Failed to retrieve resources likely due to a revoked authorization.
@@ -151,8 +151,8 @@ function DataSourceResourceSelectorChildren({
 
   const isTablesView = viewType === "tables";
   return (
-    <Tree isLoading={isResourcesLoading}>
-      {resources.map((r) => {
+    <Tree isLoading={isVaultContentLoading}>
+      {vaultContent.map((r) => {
         const isSelected = selectedResourceIds.includes(r.internalId);
         const partiallyChecked =
           !isSelected &&
@@ -186,7 +186,7 @@ function DataSourceResourceSelectorChildren({
             renderTreeItems={() => (
               <DataSourceResourceSelectorChildren
                 owner={owner}
-                dataSourceOrView={dataSourceOrView}
+                dataSourceView={dataSourceView}
                 parentId={r.internalId}
                 showExpand={showExpand}
                 // In table view, only manually selected nodes are considered and hierarchy does not apply.

diff --git a/front/components/ManagedDataSourceDocumentModal.tsx b/front/components/ManagedDataSourceDocumentModal.tsx
@@ -1,5 +1,5 @@
 import { Modal } from "@dust-tt/sparkle";
-import type { DataSourceType, WorkspaceType } from "@dust-tt/types";
+import type { DataSourceType, LightWorkspaceType } from "@dust-tt/types";
 import { useEffect, useState } from "react";
 
 export default function ManagedDataSourceDocumentModal({
@@ -9,7 +9,7 @@ export default function ManagedDataSourceDocumentModal({
   isOpen,
   setOpen,
 }: {
-  owner: WorkspaceType;
+  owner: LightWorkspaceType;
   dataSource: DataSourceType | null;
   documentId: string | null;
   isOpen: boolean;
@@ -20,7 +20,7 @@ export default function ManagedDataSourceDocumentModal({
   const [downloading, setDownloading] = useState(false);
 
   useEffect(() => {
-    if (documentId && dataSource?.name) {
+    if (documentId && dataSource) {
       setDownloading(true);
       fetch(
         `/api/w/${owner.sId}/data_sources/${encodeURIComponent(

diff --git a/front/components/assistant/AssistantDetails.tsx b/front/components/assistant/AssistantDetails.tsx
@@ -17,11 +17,11 @@ import type {
   AgentActionConfigurationType,
   AgentConfigurationScope,
   AgentConfigurationType,
-  ContentNode,
   CoreAPITable,
   DataSourceConfiguration,
   DataSourceType,
   DustAppRunConfigurationType,
+  LightContentNode,
   RetrievalConfigurationType,
   TablesQueryConfigurationType,
   WorkspaceType,
@@ -41,13 +41,13 @@ import type { KeyedMutator } from "swr";
 import { AssistantDetailsDropdownMenu } from "@app/components/assistant/AssistantDetailsDropdownMenu";
 import AssistantListActions from "@app/components/assistant/AssistantListActions";
 import { ReadOnlyTextArea } from "@app/components/assistant/ReadOnlyTextArea";
-import { SharingDropdown } from "@app/components/assistant/Sharing";
 import { assistantUsageMessage } from "@app/components/assistant/Usage";
+import { SharingDropdown } from "@app/components/assistant_builder/Sharing";
 import { PermissionTreeChildren } from "@app/components/ConnectorPermissionsTree";
 import ManagedDataSourceDocumentModal from "@app/components/ManagedDataSourceDocumentModal";
 import { SendNotificationsContext } from "@app/components/sparkle/Notification";
 import { updateAgentScope } from "@app/lib/client/dust_api";
-import { CONNECTOR_CONFIGURATIONS } from "@app/lib/connector_providers";
+import { getConnectorProviderLogo } from "@app/lib/connector_providers";
 import { getDisplayNameForDataSource } from "@app/lib/data_sources";
 import {
   useAgentConfiguration,
@@ -376,9 +376,7 @@ function DataSourcesSection({
           let dataSourceName = dsConfig.dataSourceId;
 
           if (ds) {
-            DsLogo = ds.connectorProvider
-              ? CONNECTOR_CONFIGURATIONS[ds.connectorProvider].logoComponent
-              : null;
+            DsLogo = getConnectorProviderLogo(ds.connectorProvider);
             dataSourceName = getDisplayNameForDataSource(ds);
           }
 
@@ -446,7 +444,7 @@ function DataSourceSelectedNodes({
 
   return (
     <>
-      {dataSourceSelectedNodes.nodes.map((node: ContentNode) => (
+      {dataSourceSelectedNodes.nodes.map((node: LightContentNode) => (
         <Tree.Item
           key={node.internalId}
           label={node.titleWithParentsContext ?? node.title}

diff --git a/front/components/assistant_builder/AssistantBuilder.tsx b/front/components/assistant_builder/AssistantBuilder.tsx
@@ -26,7 +26,6 @@ import { useCallback, useContext, useEffect, useMemo, useState } from "react";
 import React from "react";
 import { useSWRConfig } from "swr";
 
-import { SharingButton } from "@app/components/assistant/Sharing";
 import ActionsScreen, {
   hasActionError,
 } from "@app/components/assistant_builder/ActionsScreen";
@@ -41,6 +40,7 @@ import NamingScreen, {
   validateHandle,
 } from "@app/components/assistant_builder/NamingScreen";
 import { PrevNextButtons } from "@app/components/assistant_builder/PrevNextButtons";
+import { SharingButton } from "@app/components/assistant_builder/Sharing";
 import { submitAssistantBuilderForm } from "@app/components/assistant_builder/submitAssistantBuilderForm";
 import type {
   AssistantBuilderPendingAction,
@@ -79,12 +79,12 @@ export default function AssistantBuilder({
   baseUrl,
   defaultTemplate,
 }: AssistantBuilderProps) {
-  const { dataSources } = useContext(AssistantBuilderContext);
+  const { dataSourceViews } = useContext(AssistantBuilderContext);
   const router = useRouter();
   const { mutate } = useSWRConfig();
   const sendNotification = React.useContext(SendNotificationsContext);
-  const slackDataSource = dataSources.find(
-    (ds) => ds.connectorProvider === "slack"
+  const slackDataSource = dataSourceViews.find(
+    (dsv) => dsv.dataSource.connectorProvider === "slack"
   );
   const defaultScope =
     flow === "workspace_assistants" ? "workspace" : "private";
@@ -149,7 +149,7 @@ export default function AssistantBuilder({
     slackChannelsLinkedWithAgent,
     setSelectedSlackChannels,
   } = useSlackChannel({
-    dataSources,
+    dataSourceViews,
     initialChannels: [],
     workspaceId: owner.sId,
     isPrivateAssistant: builderState.scope === "private",
@@ -322,7 +322,7 @@ export default function AssistantBuilder({
       });
     } else {
       await mutate(
-        `/api/w/${owner.sId}/data_sources/${slackDataSource?.name}/managed/slack/channels_linked_with_agent`
+        `/api/w/${owner.sId}/data_sources/${slackDataSource?.dataSource.name}/managed/slack/channels_linked_with_agent`
       );
 
       // Redirect to the assistant list once saved.

diff --git a/front/components/assistant_builder/AssistantBuilderContext.tsx b/front/components/assistant_builder/AssistantBuilderContext.tsx
@@ -1,36 +1,27 @@
-import type {
-  AppType,
-  DataSourceType,
-  DataSourceViewType,
-} from "@dust-tt/types";
+import type { AppType, DataSourceViewType } from "@dust-tt/types";
 import { createContext } from "react";
 
 export const AssistantBuilderContext = createContext<{
   dustApps: AppType[];
-  dataSources: DataSourceType[];
   dataSourceViews: DataSourceViewType[];
 }>({
   dustApps: [],
-  dataSources: [],
   dataSourceViews: [],
 });
 
 export function AssistantBuilderProvider({
   dustApps,
-  dataSources,
   dataSourceViews,
   children,
 }: {
   dustApps: AppType[];
-  dataSources: DataSourceType[];
   dataSourceViews: DataSourceViewType[];
   children: React.ReactNode;
 }) {
   return (
     <AssistantBuilderContext.Provider
       value={{
         dustApps,
-        dataSources,
         dataSourceViews,
       }}
     >