Add dataSourceViewId in agent_data_source_configurations (#6596)

* Add dataSourceViewId in agent_data_source_configurations

* Address comments from review.

* ✨

* ✂️

* Restrict one view per data source per vault

front/lib/api/agent_data_sources.ts

@@ -13,6 +13,7 @@ import { DataSource } from "@app/lib/models/data_source";
 
 export type DataSourcesUsageByAgent = Record<ModelId, number>;
 
+// TODO(GROUPS_INFRA) Add support for views here.
 export async function getDataSourcesUsageByAgents({
   auth,
   providerFilter,

front/lib/api/assistant/configuration.ts

@@ -71,9 +71,11 @@ import {
 import { DataSource } from "@app/lib/models/data_source";
 import { Workspace } from "@app/lib/models/workspace";
 import { DataSourceResource } from "@app/lib/resources/data_source_resource";
+import { DataSourceViewResource } from "@app/lib/resources/data_source_view_resource";
 import { frontSequelize } from "@app/lib/resources/storage";
 import { generateLegacyModelSId } from "@app/lib/resources/string_ids";
 import { TemplateResource } from "@app/lib/resources/template_resource";
+import { VaultResource } from "@app/lib/resources/vault_resource";
 import { renderLightWorkspaceType } from "@app/lib/workspace";
 
 type SortStrategyType = "alphabetical" | "priority" | "updatedAt";
@@ -1230,7 +1232,7 @@ export async function createAgentActionConfiguration(
           },
           { transaction: t }
         );
-        await _createAgentDataSourcesConfigData(t, {
+        await _createAgentDataSourcesConfigData(auth, t, {
           dataSourceConfigurations: action.dataSources,
           retrievalConfigurationId: retrievalConfig.id,
           processConfigurationId: null,
@@ -1324,7 +1326,7 @@ export async function createAgentActionConfiguration(
           },
           { transaction: t }
         );
-        await _createAgentDataSourcesConfigData(t, {
+        await _createAgentDataSourcesConfigData(auth, t, {
           dataSourceConfigurations: action.dataSources,
           retrievalConfigurationId: null,
           processConfigurationId: processConfig.id,
@@ -1415,6 +1417,11 @@ function renderRetrievalTimeframeType(
   return timeframe;
 }
 
+interface DataSourceWithView {
+  ds: DataSourceResource;
+  view: DataSourceViewResource | null;
+}
+
 /**
  * Create the AgentDataSourceConfiguration rows in database.
  *
@@ -1423,6 +1430,7 @@ function renderRetrievalTimeframeType(
  * We obvisously need to do as few queries as possible.
  */
 async function _createAgentDataSourcesConfigData(
+  auth: Authenticator,
   t: Transaction,
   {
     dataSourceConfigurations,
@@ -1490,39 +1498,73 @@ async function _createAgentDataSourcesConfigData(
     []
   );
 
-  // Then we get do one findAllQuery per workspaceId, in a Promise.all
+  const globalVault = await VaultResource.fetchWorkspaceGlobalVault(auth);
+
+  // Then we get to do one findAllQuery per workspaceId, in a Promise.all.
   const getDataSourcesQueries = dsNamesPerWorkspaceId.map(
-    ({ workspace, dataSourceNames }) => {
-      return DataSourceResource.listByWorkspaceIdAndNames(
+    async ({ workspace, dataSourceNames }) => {
+      const dataSources = await DataSourceResource.listByWorkspaceIdAndNames(
         workspace,
         dataSourceNames
       );
+
+      const uniqueDataSources = _.uniqBy(dataSources, (ds) => ds.id);
+
+      // Since the UI does not currently provide the data source view,
+      // we try to retrieve the view associated with the data from the global vault
+      // and assign it to the agent data source configuration.
+      const dataSourceViews =
+        await DataSourceViewResource.listForDataSourcesInVault(
+          workspace,
+          uniqueDataSources.filter((ds) => ds.isManaged()),
+          globalVault
+        );
+
+      // Create a mapping of data source ID to data source view.
+      const dataSourceViewMap = dataSourceViews.reduce(
+        (map, view) => {
+          map[view.dataSourceId] = view;
+          return map;
+        },
+        {} as Record<string, DataSourceViewResource>
+      );
+
+      // Create an array of objects containing data sources and their corresponding data source view if it exists.
+      const dataSourcesWithViews: DataSourceWithView[] = uniqueDataSources.map(
+        (ds) => ({
+          ds,
+          view: dataSourceViewMap[ds.id] ?? null,
+        })
+      );
+
+      return dataSourcesWithViews;
     }
   );
   const results = await Promise.all(getDataSourcesQueries);
-  const dataSources = results.flat();
+  const dataSourcesWithView = results.flat();
 
   const agentDataSourcesConfigRows: AgentDataSourceConfiguration[] =
     await Promise.all(
       dataSourceConfigurations.map(async (dsConfig) => {
-        const dataSource = dataSources.find(
-          (ds) =>
-            ds.name === dsConfig.dataSourceId &&
-            ds.workspaceId ===
+        const dataSourceWithView = dataSourcesWithView.find(
+          (d) =>
+            d.ds.name === dsConfig.dataSourceId &&
+            d.ds.workspaceId ===
               workspaces.find((w) => w.sId === dsConfig.workspaceId)?.id
         );
-        if (!dataSource) {
+        if (!dataSourceWithView) {
           throw new Error(
             `Can't create AgentDataSourcesConfig: datasource not found. dataSourceId: ${dsConfig.dataSourceId}`
           );
         }
         return AgentDataSourceConfiguration.create(
           {
-            dataSourceId: dataSource.id,
+            dataSourceId: dataSourceWithView.ds.id,
             parentsIn: dsConfig.filter.parents?.in,
             parentsNotIn: dsConfig.filter.parents?.not,
             retrievalConfigurationId: retrievalConfigurationId,
             processConfigurationId: processConfigurationId,
+            dataSourceViewId: dataSourceWithView.view?.id,
           },
           { transaction: t }
         );

front/lib/models/assistant/actions/data_sources.ts

@@ -11,6 +11,7 @@ import { AgentProcessConfiguration } from "@app/lib/models/assistant/actions/pro
 import { AgentRetrievalConfiguration } from "@app/lib/models/assistant/actions/retrieval";
 import { DataSource } from "@app/lib/models/data_source";
 import { frontSequelize } from "@app/lib/resources/storage";
+import { DataSourceViewModel } from "@app/lib/resources/storage/models/data_source_view";
 
 /**
  * Configuration of Datasources used for Retrieval Action.
@@ -27,6 +28,7 @@ export class AgentDataSourceConfiguration extends Model<
   declare parentsNotIn: string[] | null;
 
   declare dataSourceId: ForeignKey<DataSource["id"]>;
+  declare dataSourceViewId: ForeignKey<DataSourceViewModel["id"]>;
 
   // AgentDataSourceConfiguration can be used by both the retrieval and the process actions'
   // configurations.
@@ -114,3 +116,12 @@ AgentDataSourceConfiguration.belongsTo(DataSource, {
   as: "dataSource",
   foreignKey: { name: "dataSourceId", allowNull: false },
 });
+
+// Data source config <> Data source view
+DataSourceViewModel.hasMany(AgentDataSourceConfiguration, {
+  foreignKey: { allowNull: true },
+  onDelete: "RESTRICT",
+});
+AgentDataSourceConfiguration.belongsTo(DataSourceViewModel, {
+  foreignKey: { allowNull: false },
+});

front/lib/resources/data_source_view_resource.ts

@@ -4,6 +4,7 @@
 import type {
   DataSourceType,
   DataSourceViewType,
+  LightWorkspaceType,
   ModelId,
   Result,
 } from "@dust-tt/types";
@@ -83,6 +84,24 @@ export class DataSourceViewResource extends BaseResource<DataSourceViewModel> {
     return blobs.map((b) => new this(this.model, b.get()));
   }
 
+  static async listForDataSourcesInVault(
+    owner: LightWorkspaceType,
+    dataSources: DataSourceResource[],
+    vault: VaultResource,
+    { transaction }: { transaction?: Transaction } = {}
+  ) {
+    const blobs = await this.model.findAll({
+      where: {
+        workspaceId: owner.id,
+        dataSourceId: dataSources.map((ds) => ds.id),
+        vaultId: vault.id,
+      },
+      transaction,
+    });
+
+    return blobs.map((b) => new this(this.model, b.get()));
+  }
+
   static async fetchById(auth: Authenticator, id: string) {
     const fileModelId = getResourceIdFromSId(id);
     if (!fileModelId) {

front/lib/resources/storage/models/data_source_view.ts

@@ -53,7 +53,7 @@ DataSourceViewModel.init(
     indexes: [
       { fields: ["workspaceId", "id"] },
       { fields: ["workspaceId", "vaultId"] },
-      { fields: ["workspaceId", "dataSourceId"] },
+      { fields: ["workspaceId", "dataSourceId", "vaultId"], unique: true },
     ],
   }
 );

front/migrations/db/migration_51.sql

@@ -0,0 +1,5 @@
+-- Migration created on Jul 30, 2024
+CREATE UNIQUE INDEX "data_source_views_workspace_id_data_source_id_vault_id" ON "data_source_views" ("workspaceId", "dataSourceId", "vaultId");
+
+ALTER TABLE "public"."agent_data_source_configurations"
+ADD COLUMN "dataSourceViewId" INTEGER REFERENCES "data_source_views" ("id") ON DELETE RESTRICT ON UPDATE CASCADE;