fix(doc tracker revival): use vaults (#8143)

Co-authored-by: Henry Fontanier <[email protected]>

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

@@ -1,43 +1,21 @@
-import { DustAPI } from "@dust-tt/client";
-import type { WorkspaceType } from "@dust-tt/types";
-
-import config from "@app/lib/api/config";
-import { prodAPICredentialsForOwner } from "@app/lib/auth";
+import type { Authenticator } from "@app/lib/auth";
 import { TRACKABLE_CONNECTOR_TYPES } from "@app/lib/documents_post_process_hooks/hooks/document_tracker/consts";
-import logger from "@app/logger/logger";
-
-export async function getTrackableDataSources(owner: WorkspaceType): Promise<
-  {
-    workspace_id: string;
-    data_source_id: string;
-  }[]
-> {
-  const prodCredentials = await prodAPICredentialsForOwner(owner);
-
-  const prodAPI = new DustAPI(
-    config.getDustAPIConfig(),
-    prodCredentials,
-    logger
-  );
+import { DataSourceViewResource } from "@app/lib/resources/data_source_view_resource";
+import { VaultResource } from "@app/lib/resources/vault_resource";
 
-  // Fetch data sources
-  const dsRes = await prodAPI.getDataSources(prodAPI.workspaceId());
-  if (dsRes.isErr()) {
-    throw dsRes.error;
-  }
-  const dataSources = dsRes.value;
+export async function getTrackableDataSourceViews(
+  auth: Authenticator
+): Promise<DataSourceViewResource[]> {
+  const globalVault = await VaultResource.fetchWorkspaceGlobalVault(auth);
+  // TODO(DOC_TRACKER):
+  const views = await DataSourceViewResource.listByVault(auth, globalVault);
 
   // Filter data sources to only include trackable ones
-  const trackableDataSources = dataSources.filter(
-    (ds) =>
-      ds.connectorProvider &&
-      TRACKABLE_CONNECTOR_TYPES.includes(ds.connectorProvider)
+  const trackableViews = views.filter(
+    (view) =>
+      view.dataSource.connectorProvider &&
+      TRACKABLE_CONNECTOR_TYPES.includes(view.dataSource.connectorProvider)
   );
 
-  return trackableDataSources.map((ds) => ({
-    workspace_id: prodAPI.workspaceId(),
-    // TODO(GROUPS_INFRA): this should pull the data source views from the global vaults (we need an
-    // API for that).
-    data_source_id: ds.sId,
-  }));
+  return trackableViews;
 }

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

@@ -2,7 +2,7 @@ import * as t from "io-ts";
 
 import { callAction } from "@app/lib/actions/helpers";
 import type { Authenticator } from "@app/lib/auth";
-import { getTrackableDataSources } from "@app/lib/documents_post_process_hooks/hooks/document_tracker/lib";
+import { getTrackableDataSourceViews } from "@app/lib/documents_post_process_hooks/hooks/document_tracker/lib";
 import { cloneBaseConfig, DustProdActionRegistry } from "@app/lib/registry";
 
 // Part of the new doc tracker pipeline, performs the retrieval (semantic search) step
@@ -16,8 +16,17 @@ export async function callDocTrackerRetrievalAction(
   const action = DustProdActionRegistry["doc-tracker-retrieval"];
   const config = cloneBaseConfig(action.config);
 
-  config.SEMANTIC_SEARCH.data_sources = await getTrackableDataSources(
-    auth.getNonNullableWorkspace()
+  const trackableDataSourceViews = await getTrackableDataSourceViews(auth);
+
+  if (!trackableDataSourceViews.length) {
+    return [];
+  }
+
+  config.SEMANTIC_SEARCH.data_sources = trackableDataSourceViews.map(
+    (view) => ({
+      workspace_id: auth.getNonNullableWorkspace().sId,
+      data_source_id: view.sId,
+    })
   );
   config.SEMANTIC_SEARCH.target_document_tokens = targetDocumentTokens;