Upg: show usage by agents in all screens. Opt: improve performance of…

… usage query. Fix: handle all cases of usages (table, process, retrieval).
dust-tt · Sep 18, 2024 · 7e1367b · 7e1367b
1 parent bcc6ebd
commit 7e1367b
Show file tree

Hide file tree

Showing 23 changed files with 1,011 additions and 188 deletions.
diff --git a/front/components/data_source/DeleteStaticDataSourceDialog.tsx b/front/components/data_source/DeleteStaticDataSourceDialog.tsx
@@ -1,13 +1,13 @@
 import { Dialog } from "@dust-tt/sparkle";
-import type { DataSourceType } from "@dust-tt/types";
+import type { DataSourceType, DataSourceUsageType } from "@dust-tt/types";
 import { useState } from "react";
 
 import { getDataSourceName, isManaged } from "@app/lib/data_sources";
 
 interface DeleteStaticDataSourceDialogProps {
   dataSource: DataSourceType;
   handleDelete: () => void;
-  dataSourceUsage?: number;
+  dataSourceUsage?: DataSourceUsageType;
   isOpen: boolean;
   onClose: () => void;
 }
@@ -34,8 +34,8 @@ export function DeleteStaticDataSourceDialog({
   const message =
     dataSourceUsage === undefined
       ? `Are you sure you want to permanently delete ${name}?`
-      : dataSourceUsage > 0
-        ? `${dataSourceUsage} assistants currently use ${name}. Are you sure you want to remove?`
+      : dataSourceUsage.count > 0
+        ? `${dataSourceUsage} assistants currently use ${name}: ${dataSourceUsage.agentNames.join(", ")}. Are you sure you want to remove?`
         : `No assistants are using ${name}. Confirm permanent deletion?`;
 
   return (

diff --git a/front/components/data_source/WebsiteConfiguration.tsx b/front/components/data_source/WebsiteConfiguration.tsx
@@ -12,6 +12,7 @@ import {
 import type {
   CrawlingFrequency,
   DataSourceType,
+  DataSourceUsageType,
   DepthOption,
   SubscriptionType,
   UpdateConnectorConfigurationType,
@@ -50,7 +51,7 @@ export default function WebsiteConfiguration({
   dataSources: DataSourceType[];
   webCrawlerConfiguration: WebCrawlerConfigurationType | null;
   dataSource: DataSourceType | null;
-  dataSourceUsage?: number;
+  dataSourceUsage?: DataSourceUsageType;
 }) {
   const [isSaving, setIsSaving] = useState(false);
   const [isSubmitted, setIsSubmitted] = useState(false);
@@ -535,7 +536,7 @@ export default function WebsiteConfiguration({
                   handleDelete={handleDelete}
                   isOpen={isDeleteModalOpen}
                   onClose={() => setIsDeleteModalOpen(false)}
-                  dataSourceUsage={dataSourceUsage ?? 0}
+                  dataSourceUsage={dataSourceUsage}
                 />
               )}
             </div>

diff --git a/front/components/vaults/CreateOrEditVaultModal.tsx b/front/components/vaults/CreateOrEditVaultModal.tsx
@@ -300,6 +300,7 @@ export function CreateOrEditVaultModal({
                   size="sm"
                   label="Delete Vault"
                   variant="primaryWarning"
+                  className="mr-2"
                   onClick={() => setShowDeleteConfirmDialog(true)}
                 />
               </div>

diff --git a/front/components/vaults/EditVaultStaticDatasourcesViews.tsx b/front/components/vaults/EditVaultStaticDatasourcesViews.tsx
@@ -1,7 +1,7 @@
 import { Button, PlusIcon, Popup, Tooltip } from "@dust-tt/sparkle";
 import type {
   DataSourceType,
-  DataSourceViewWithConnectorType,
+  DataSourceViewType,
   PlanType,
   VaultType,
   WorkspaceType,
@@ -22,7 +22,7 @@ interface EditVaultStaticDatasourcesViewsProps {
   plan: PlanType;
   vault: VaultType;
   dataSources: DataSourceType[];
-  dataSourceView: DataSourceViewWithConnectorType | null;
+  dataSourceView: DataSourceViewType | null;
   category: "folder" | "website";
   onClose: () => void;
 }

diff --git a/front/components/vaults/VaultCategoriesList.tsx b/front/components/vaults/VaultCategoriesList.tsx
@@ -10,7 +10,11 @@ import {
   Searchbar,
   Spinner,
 } from "@dust-tt/sparkle";
-import type { VaultType, WorkspaceType } from "@dust-tt/types";
+import type {
+  DataSourceUsageType,
+  VaultType,
+  WorkspaceType,
+} from "@dust-tt/types";
 import { DATA_SOURCE_VIEW_CATEGORIES, removeNulls } from "@dust-tt/types";
 import type { CellContext } from "@tanstack/react-table";
 import type { ComponentType } from "react";
@@ -23,7 +27,7 @@ type RowData = {
   category: string;
   name: string;
   icon: ComponentType;
-  usage: number;
+  usage: DataSourceUsageType;
   count: number;
   onClick?: () => void;
 };
@@ -84,8 +88,11 @@ const getTableColumns = () => {
       cell: (info: Info) => (
         <>
           {info.row.original.usage ? (
-            <DataTable.CellContent icon={RobotIcon}>
-              {info.row.original.usage}
+            <DataTable.CellContent
+              icon={RobotIcon}
+              title={`Used by ${info.row.original.usage.agentNames.join(", ")}`}
+            >
+              {info.row.original.usage.count}
             </DataTable.CellContent>
           ) : null}
         </>

diff --git a/front/components/vaults/VaultDataSourceViewContentList.tsx b/front/components/vaults/VaultDataSourceViewContentList.tsx
@@ -91,7 +91,7 @@ const getTableColumns = (showVaultUsage: boolean): ColumnDef<RowData>[] => {
             ? info
                 .getValue()
                 .map((v) => v.name)
-                .join(",")
+                .join(", ")
             : "-"}
         </DataTable.CellContent>
       ),

diff --git a/front/components/vaults/VaultResourcesList.tsx b/front/components/vaults/VaultResourcesList.tsx
@@ -6,22 +6,25 @@ import {
   DataTable,
   FolderIcon,
   PencilSquareIcon,
+  RobotIcon,
   Searchbar,
   Spinner,
   TrashIcon,
   usePaginationFromUrl,
 } from "@dust-tt/sparkle";
 import type {
   ConnectorProvider,
+  DataSourceUsageType,
   DataSourceViewCategory,
+  DataSourceViewsWithDetails,
   DataSourceViewType,
-  DataSourceViewWithConnectorType,
+  DataSourceWithConnectorDetailsType,
   PlanType,
   VaultType,
   WorkspaceType,
 } from "@dust-tt/types";
 import { isWebsiteOrFolderCategory } from "@dust-tt/types";
-import type { CellContext, ColumnDef } from "@tanstack/react-table";
+import type { CellContext, ColumnDef, Row } from "@tanstack/react-table";
 import { useRouter } from "next/router";
 import type { ComponentType } from "react";
 import { useMemo } from "react";
@@ -41,7 +44,7 @@ import { getDataSourceNameFromView } from "@app/lib/data_sources";
 import { useDataSources } from "@app/lib/swr/data_sources";
 import {
   useDeleteFolderOrWebsite,
-  useVaultDataSourceViews,
+  useVaultDataSourceViewsWithDetails,
 } from "@app/lib/swr/vaults";
 import { classNames } from "@app/lib/utils";
 
@@ -54,7 +57,7 @@ const REDIRECT_TO_EDIT_PERMISSIONS = [
 ];
 
 type RowData = {
-  dataSourceView: DataSourceViewWithConnectorType;
+  dataSourceView: DataSourceViewsWithDetails;
   label: string;
   icon: ComponentType;
   workspaceId: string;
@@ -121,6 +124,35 @@ const getTableColumns = ({
     },
   };
 
+  const usedByColumn = {
+    header: "Used by",
+    id: "usedBy",
+    accessorFn: (row: RowData) => row.dataSourceView.usage,
+    sortingFn: (rowA: Row<RowData>, rowB: Row<RowData>) => {
+      return (
+        (rowA.original.dataSourceView.usage?.count ?? 0) -
+        (rowB.original.dataSourceView.usage?.count ?? 0)
+      );
+    },
+    meta: {
+      width: "6rem",
+    },
+    cell: (info: CellContext<RowData, DataSourceUsageType>) => (
+      <>
+        {info.row.original.dataSourceView.usage ? (
+          <DataTable.CellContent
+            icon={RobotIcon}
+            title={`Used by ${info.row.original.dataSourceView.usage.agentNames.join(
+              ", "
+            )}`}
+          >
+            {info.row.original.dataSourceView.usage.count}
+          </DataTable.CellContent>
+        ) : null}
+      </>
+    ),
+  };
+
   const lastSyncedColumn = {
     header: "Last sync",
     id: "lastSync",
@@ -197,11 +229,17 @@ const getTableColumns = ({
 
   // TODO(GROUPS_UI) Add usage column.
   if (isSystemVault && isManaged) {
-    return [nameColumn, managedByColumn, lastSyncedColumn, actionColumn];
+    return [
+      nameColumn,
+      usedByColumn,
+      managedByColumn,
+      lastSyncedColumn,
+      actionColumn,
+    ];
   }
   return isManaged
-    ? [nameColumn, managedByColumn, lastSyncedColumn]
-    : [nameColumn, managedByColumn];
+    ? [nameColumn, usedByColumn, managedByColumn, lastSyncedColumn]
+    : [nameColumn, usedByColumn, managedByColumn];
 };
 
 export const VaultResourcesList = ({
@@ -220,7 +258,7 @@ export const VaultResourcesList = ({
   const [showConnectorPermissionsModal, setShowConnectorPermissionsModal] =
     useState(false);
   const [selectedDataSourceView, setSelectedDataSourceView] =
-    useState<DataSourceViewWithConnectorType | null>(null);
+    useState<DataSourceViewsWithDetails | null>(null);
   const [showDeleteConfirmDialog, setShowDeleteConfirmDialog] = useState(false);
   const [showFolderOrWebsiteModal, setShowFolderOrWebsiteModal] =
     useState(false);
@@ -253,12 +291,10 @@ export const VaultResourcesList = ({
     vaultDataSourceViews,
     isVaultDataSourceViewsLoading,
     mutateRegardlessOfQueryParams: mutateVaultDataSourceViews,
-  } = useVaultDataSourceViews({
+  } = useVaultDataSourceViewsWithDetails({
     workspaceId: owner.sId,
     vaultId: vault.sId,
     category: category,
-    includeEditedBy: true,
-    includeConnectorDetails: true,
   });
 
   const rows: RowData[] = useMemo(
@@ -294,7 +330,7 @@ export const VaultResourcesList = ({
           icon: getConnectorProviderLogoWithFallback(provider, FolderIcon),
           workspaceId: owner.sId,
           isAdmin,
-          isLoading: isLoadingByProvider[provider],
+          isLoading: provider ? isLoadingByProvider[provider] : false,
           moreMenuItems,
           buttonOnClick: (e) => {
             e.stopPropagation();
@@ -367,9 +403,16 @@ export const VaultResourcesList = ({
                 owner={owner}
                 dustClientFacingUrl={dustClientFacingUrl}
                 plan={plan}
-                existingDataSources={vaultDataSourceViews.map(
-                  (v) => v.dataSource
-                )}
+                existingDataSources={
+                  vaultDataSourceViews
+                    .filter(
+                      (dsView) => dsView.dataSource.connectorProvider !== null
+                    )
+                    .map(
+                      (v) => v.dataSource
+                    ) as DataSourceWithConnectorDetailsType[]
+                  // We need to filter and then cast because useVaultDataSourceViewsWithDetails can return dataSources with connectorProvider as null
+                }
                 setIsProviderLoading={(provider, isLoading) => {
                   setIsNewConnectorLoading(isLoading);
                   setIsLoadingByProvider((prev) => ({