Flav/views only for managed datasources (#6608)

* s/datasource/data_source

* Only create a data source view for managed data sources with a connection

* Create data source with provider

* Add migration to delete un managed data source views

* 🙈

* ✨

* ✨

front/lib/api/assistant/configuration.ts

@@ -70,7 +70,7 @@ import {
 } from "@app/lib/models/assistant/conversation";
 import { DataSource } from "@app/lib/models/data_source";
 import { Workspace } from "@app/lib/models/workspace";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import { frontSequelize } from "@app/lib/resources/storage";
 import { generateLegacyModelSId } from "@app/lib/resources/string_ids";
 import { TemplateResource } from "@app/lib/resources/template_resource";

front/lib/api/data_sources.ts

@@ -10,7 +10,7 @@ import config from "@app/lib/api/config";
 import { getMembers } from "@app/lib/api/workspace";
 import type { Authenticator } from "@app/lib/auth";
 import { sendGithubDeletionEmail } from "@app/lib/email";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import logger from "@app/logger/logger";
 import { launchScrubDataSourceWorkflow } from "@app/poke/temporal/client";
 

front/lib/data_sources.ts

@@ -2,7 +2,6 @@ import type { CoreAPIDocument, DataSourceType } from "@dust-tt/types";
 import { assertNever } from "@dust-tt/types";
 
 import { CONNECTOR_CONFIGURATIONS } from "@app/lib/connector_providers";
-import type { DataSource } from "@app/lib/models/data_source";
 
 export function getDisplayNameForDocument(document: CoreAPIDocument): string {
   const titleTagPrefix = "title:";
@@ -34,17 +33,3 @@ export function getDisplayNameForDataSource(ds: DataSourceType) {
     return ds.name;
   }
 }
-
-export function renderDataSourceType(dataSource: DataSource): DataSourceType {
-  return {
-    id: dataSource.id,
-    createdAt: dataSource.createdAt.getTime(),
-    name: dataSource.name,
-    description: dataSource.description,
-    assistantDefaultSelected: dataSource.assistantDefaultSelected,
-    dustAPIProjectId: dataSource.dustAPIProjectId,
-    connectorId: dataSource.connectorId,
-    connectorProvider: dataSource.connectorProvider,
-    editedByUser: undefined,
-  };
-}

front/lib/document_tracker.ts

@@ -5,7 +5,7 @@ import { Authenticator } from "@app/lib/auth";
 import { TrackedDocument } from "@app/lib/models/doc_tracker";
 import { User } from "@app/lib/models/user";
 import { Workspace } from "@app/lib/models/workspace";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import { MembershipResource } from "@app/lib/resources/membership_resource";
 import logger from "@app/logger/logger";
 

front/lib/documents_post_process_hooks/hooks/data_source_helpers.ts

@@ -5,7 +5,7 @@ import config from "@app/lib/api/config";
 import type { Authenticator } from "@app/lib/auth";
 import { diffStrings } from "@app/lib/diff";
 import { Workspace } from "@app/lib/models/workspace";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import logger from "@app/logger/logger";
 
 export async function getPreviousDocumentVersion({

front/lib/documents_post_process_hooks/hooks/document_tracker/suggest_changes/lib.ts

@@ -17,7 +17,7 @@ import {
   TrackedDocument,
 } from "@app/lib/models/doc_tracker";
 import { User } from "@app/lib/models/user";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import mainLogger from "@app/logger/logger";
 
 import { callDocTrackerRetrievalAction } from "./actions/doc_tracker_retrieval";

front/lib/resources/datasource_resource.ts → front/lib/resources/data_source_resource.ts

@@ -163,7 +163,9 @@ export class DataSourceResource extends BaseResource<DataSource> {
     auth: Authenticator,
     transaction?: Transaction
   ): Promise<Result<undefined, Error>> {
-    await DataSourceViewResource.deleteForDataSource(auth, this);
+    if (this.isManaged()) {
+      await DataSourceViewResource.deleteForDataSource(auth, this);
+    }
 
     try {
       await this.model.destroy({
@@ -212,6 +214,16 @@ export class DataSourceResource extends BaseResource<DataSource> {
     };
   }
 
+  isManaged(): boolean {
+    return (
+      this.name.startsWith("managed-") &&
+      this.connectorProvider !== null &&
+      this.connectorProvider !== "webcrawler"
+    );
+  }
+
+  // Serialization.
+
   toJSON(): DataSourceType {
     return {
       id: this.id,

front/lib/resources/data_source_view_resource.ts

@@ -17,7 +17,7 @@ import type {
 
 import type { Authenticator } from "@app/lib/auth";
 import { BaseResource } from "@app/lib/resources/base_resource";
-import type { DataSourceResource } from "@app/lib/resources/datasource_resource";
+import type { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import { DataSourceViewModel } from "@app/lib/resources/storage/models/data_source_view";
 import type { ReadonlyAttributesType } from "@app/lib/resources/storage/types";
 import { getResourceIdFromSId, makeSId } from "@app/lib/resources/string_ids";
@@ -57,12 +57,11 @@ export class DataSourceViewResource extends BaseResource<DataSourceViewModel> {
     });
   }
 
-  // TODO(2024-07-29 flav) Replace dataSourceId by DataSourceResource once implemented.
   // For now, we create a default view for all data sources in the global vault.
   // This view has access to all documents, which is represented by null.
   static async createViewInVaultFromDataSourceIncludingAllDocuments(
     vault: VaultResource,
-    dataSource: DataSourceType
+    dataSource: DataSourceResource
   ) {
     return this.makeNew({
       dataSourceId: dataSource.id,

front/lib/workspace_usage.ts

@@ -11,7 +11,7 @@ import {
 } from "@app/lib/models/assistant/conversation";
 import { User } from "@app/lib/models/user";
 import { Workspace } from "@app/lib/models/workspace";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 
 import { frontSequelize } from "./resources/storage";
 
@@ -185,7 +185,7 @@ export async function getMessageUsageData(
         "workspaces" w ON c."workspaceId" = w."id"
       LEFT JOIN
         "agent_configurations" ac ON am."agentConfigurationId" = ac."sId" AND am."agentConfigurationVersion" = ac."version"
-      LEFT JOIN 
+      LEFT JOIN
         "messages" m2 on m."parentId" = m2."id"
       LEFT JOIN
         "user_messages" um on m2."userMessageId" = um."id"

front/migrations/20240730_backfill_data_source_views.ts

@@ -1,8 +1,7 @@
 import type { LightWorkspaceType } from "@dust-tt/types";
 
 import { Authenticator } from "@app/lib/auth";
-import { renderDataSourceType } from "@app/lib/data_sources";
-import { DataSource } from "@app/lib/models/data_source";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import { DataSourceViewResource } from "@app/lib/resources/data_source_view_resource";
 import { DataSourceViewModel } from "@app/lib/resources/storage/models/data_source_view";
 import { VaultResource } from "@app/lib/resources/vault_resource";
@@ -14,11 +13,8 @@ async function backfillDataSourceViewsForWorkspace(
   logger: Logger,
   execute: boolean
 ) {
-  const dataSources = await DataSource.findAll({
-    where: {
-      workspaceId: workspace.id,
-    },
-  });
+  const auth = await Authenticator.internalAdminForWorkspace(workspace.sId);
+  const dataSources = await DataSourceResource.listByWorkspace(auth);
 
   logger.info(
     `Found ${dataSources.length} data sources for workspace(${workspace.sId}).`
@@ -28,8 +24,6 @@ async function backfillDataSourceViewsForWorkspace(
     return;
   }
 
-  const auth = await Authenticator.internalAdminForWorkspace(workspace.sId);
-
   const globalVault = await VaultResource.fetchWorkspaceGlobalVault(auth);
 
   let updated = 0;
@@ -53,7 +47,7 @@ async function backfillDataSourceViewsForWorkspace(
     // Create a view for this data source in the global vault.
     await DataSourceViewResource.createViewInVaultFromDataSourceIncludingAllDocuments(
       globalVault,
-      renderDataSourceType(dataSource)
+      dataSource
     );
 
     updated++;

front/migrations/20240730_delete_unmanaged_data_source_views.ts

@@ -0,0 +1,64 @@
+import type { LightWorkspaceType } from "@dust-tt/types";
+import assert from "assert";
+
+import { Authenticator } from "@app/lib/auth";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
+import { DataSourceViewResource } from "@app/lib/resources/data_source_view_resource";
+import type { Logger } from "@app/logger/logger";
+import { makeScript, runOnAllWorkspaces } from "@app/scripts/helpers";
+
+async function deleteUnmanagedDataSourceViewsForWorkspace(
+  workspace: LightWorkspaceType,
+  logger: Logger,
+  execute: boolean
+) {
+  const auth = await Authenticator.internalAdminForWorkspace(workspace.sId);
+
+  // List workspace data sources.
+  const dataSources = await DataSourceResource.listByWorkspace(auth);
+
+  // List workspace data source views.
+  const dataSourceViews = await DataSourceViewResource.listByWorkspace(auth);
+
+  logger.info(
+    `Found ${dataSourceViews.length} data source views for workspace(${workspace.sId}).`
+  );
+
+  const viewsToDelete: DataSourceViewResource[] = [];
+  for (const dataSourceView of dataSourceViews) {
+    const ds = dataSources.find((ds) => ds.id === dataSourceView.dataSourceId);
+    assert(ds, `Data source ${dataSourceView.dataSourceId} not found.`);
+
+    // If data source is not managed, delete the view.
+    if (!ds.isManaged()) {
+      viewsToDelete.push(dataSourceView);
+    }
+  }
+
+  logger.info(
+    `About to delete ${viewsToDelete.length} unmanaged data source views for workspace(${workspace.sId}).`
+  );
+
+  if (!execute) {
+    return;
+  }
+
+  for (const view of viewsToDelete) {
+    await view.delete(auth);
+    logger.info(`Deleted view for data source ${view.dataSourceId}.`);
+  }
+
+  logger.info(
+    `Deleted all unmanaged data source views for workspace(${workspace.sId}).`
+  );
+}
+
+makeScript({}, async ({ execute }, logger) => {
+  return runOnAllWorkspaces(async (workspace) => {
+    await deleteUnmanagedDataSourceViewsForWorkspace(
+      workspace,
+      logger,
+      execute
+    );
+  });
+});

front/pages/api/poke/workspaces/[wId]/data_sources/[name]/config.ts

@@ -5,7 +5,7 @@ import type { NextApiRequest, NextApiResponse } from "next";
 import config from "@app/lib/api/config";
 import { withSessionAuthentication } from "@app/lib/api/wrappers";
 import { Authenticator, getSession } from "@app/lib/auth";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import logger from "@app/logger/logger";
 import { apiError } from "@app/logger/withlogging";
 

front/pages/api/registry/[type]/lookup.ts

@@ -3,7 +3,7 @@ import type { NextApiRequest, NextApiResponse } from "next";
 
 import { Authenticator } from "@app/lib/auth";
 import { Workspace } from "@app/lib/models/workspace";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import { withLogging } from "@app/logger/withlogging";
 
 const { DUST_REGISTRY_SECRET } = process.env;

front/pages/api/w/[wId]/assistant/agent_configurations/[aId]/linked_slack_channels.ts

@@ -9,7 +9,7 @@ import { getAgentConfiguration } from "@app/lib/api/assistant/configuration";
 import config from "@app/lib/api/config";
 import { withSessionAuthenticationForWorkspace } from "@app/lib/api/wrappers";
 import type { Authenticator } from "@app/lib/auth";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import logger from "@app/logger/logger";
 import { apiError } from "@app/logger/withlogging";
 

front/pages/api/w/[wId]/data_sources/[name]/index.ts

@@ -7,7 +7,7 @@ import {
 } from "@app/lib/api/data_sources";
 import { withSessionAuthenticationForWorkspace } from "@app/lib/api/wrappers";
 import type { Authenticator } from "@app/lib/auth";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import { apiError } from "@app/logger/withlogging";
 
 export type GetOrPostDataSourceResponseBody = {

front/pages/api/w/[wId]/data_sources/index.ts

@@ -12,8 +12,7 @@ import config from "@app/lib/api/config";
 import { getDataSource, getDataSources } from "@app/lib/api/data_sources";
 import { withSessionAuthenticationForWorkspace } from "@app/lib/api/wrappers";
 import type { Authenticator } from "@app/lib/auth";
-import { DataSourceViewResource } from "@app/lib/resources/data_source_view_resource";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import { VaultResource } from "@app/lib/resources/vault_resource";
 import { ServerSideTracking } from "@app/lib/tracking/server";
 import logger from "@app/logger/logger";
@@ -183,11 +182,6 @@ async function handler(
         vaultId: globalVault.id,
       });
 
-      await DataSourceViewResource.createViewInVaultFromDataSourceIncludingAllDocuments(
-        globalVault,
-        ds.toJSON()
-      );
-
       const dataSourceType = await getDataSource(auth, ds.name);
       if (dataSourceType) {
         void ServerSideTracking.trackDataSourceCreated({

front/pages/api/w/[wId]/data_sources/managed.ts

@@ -23,8 +23,8 @@ import { getDataSource } from "@app/lib/api/data_sources";
 import { withSessionAuthenticationForWorkspace } from "@app/lib/api/wrappers";
 import type { Authenticator } from "@app/lib/auth";
 import { getOrCreateSystemApiKey } from "@app/lib/auth";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import { DataSourceViewResource } from "@app/lib/resources/data_source_view_resource";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
 import { VaultResource } from "@app/lib/resources/vault_resource";
 import { ServerSideTracking } from "@app/lib/tracking/server";
 import { isDisposableEmailDomain } from "@app/lib/utils/disposable_email_domains";
@@ -332,24 +332,25 @@ async function handler(
         ? VaultResource.fetchWorkspaceGlobalVault(auth)
         : VaultResource.fetchWorkspaceSystemVault(auth));
       const dataSource = await DataSourceResource.makeNew({
-        name: dataSourceName,
+        assistantDefaultSelected,
+        connectorProvider: provider,
         description: dataSourceDescription,
         dustAPIProjectId: dustProject.value.project.project_id.toString(),
-        workspaceId: owner.id,
-        assistantDefaultSelected,
         editedByUserId: user.id,
+        name: dataSourceName,
         vaultId: vault.id,
+        workspaceId: owner.id,
       });
 
-      const globalVault = vault.isGlobal()
-        ? vault
-        : await VaultResource.fetchWorkspaceGlobalVault(auth);
-
       // For managed data source, we create a default view in the workspace vault.
-      await DataSourceViewResource.createViewInVaultFromDataSourceIncludingAllDocuments(
-        globalVault,
-        dataSource.toJSON()
-      );
+      if (dataSource.isManaged()) {
+        const globalVault = await VaultResource.fetchWorkspaceGlobalVault(auth);
+
+        await DataSourceViewResource.createViewInVaultFromDataSourceIncludingAllDocuments(
+          globalVault,
+          dataSource
+        );
+      }
 
       const connectorsAPI = new ConnectorsAPI(
         config.getConnectorsAPIConfig(),
@@ -410,7 +411,6 @@ async function handler(
 
       await dataSource.update({
         connectorId: connectorsRes.value.id,
-        connectorProvider: provider,
       });
       const dataSourceType = await getDataSource(auth, dataSource.name);
       if (dataSourceType) {

front/poke/temporal/activities.ts

@@ -44,8 +44,8 @@ import { Subscription } from "@app/lib/models/plan";
 import { UserMetadata } from "@app/lib/models/user";
 import { MembershipInvitation, Workspace } from "@app/lib/models/workspace";
 import { ContentFragmentResource } from "@app/lib/resources/content_fragment_resource";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import { DataSourceViewResource } from "@app/lib/resources/data_source_view_resource";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
 import { FileResource } from "@app/lib/resources/file_resource";
 import { GroupResource } from "@app/lib/resources/group_resource";
 import { KeyResource } from "@app/lib/resources/key_resource";

front/temporal/documents_post_process_hooks/activities.ts

@@ -11,7 +11,7 @@ import { Authenticator } from "@app/lib/auth";
 import type { DocumentsPostProcessHookType } from "@app/lib/documents_post_process_hooks/hooks";
 import { DOCUMENTS_POST_PROCESS_HOOK_BY_TYPE } from "@app/lib/documents_post_process_hooks/hooks";
 import { Workspace } from "@app/lib/models/workspace";
-import { DataSourceResource } from "@app/lib/resources/datasource_resource";
+import { DataSourceResource } from "@app/lib/resources/data_source_resource";
 import { withRetries } from "@app/lib/utils/retries";
 import logger from "@app/logger/logger";