feat: add link to maintained document in tracker email

dust-tt · Jan 13, 2025 · 87d36cf · 87d36cf
1 parent 4f51273
commit 87d36cf
Show file tree

Hide file tree

Showing 4 changed files with 103 additions and 36 deletions.
diff --git a/front/lib/api/tracker.ts b/front/lib/api/tracker.ts
@@ -1,5 +1,5 @@
 import type { TrackerGenerationToProcess } from "@dust-tt/types";
-import { concurrentExecutor, CoreAPI } from "@dust-tt/types";
+import { concurrentExecutor, CoreAPI, removeNulls } from "@dust-tt/types";
 import _ from "lodash";
 
 import config from "@app/lib/api/config";
@@ -127,15 +127,32 @@ const _sendTrackerWithGenerationEmail = async ({
   localLogger: Logger;
 }): Promise<void> => {
   const coreAPI = new CoreAPI(config.getCoreAPIConfig(), localLogger);
-  const generationsByDataSources = _.groupBy(generations, "dataSource.id");
-  const documentsById = new Map();
+  const dataSourceById = _.keyBy(
+    removeNulls(
+      generations.map((g) => [g.dataSource, g.maintainedDataSource]).flat()
+    ),
+    "id"
+  );
+  const docsToFetchByDataSourceId = _.mapValues(
+    _.groupBy(
+      generations.map((g) => ({
+        dataSourceId: g.dataSource.id,
+        documentIds: removeNulls([g.documentId, g.maintainedDocumentId]),
+      })),
+      "dataSourceId"
+    ),
+    (docs) => docs.map((d) => d.documentIds).flat()
+  );
+  const documentsByIdentifier = new Map<
+    string,
+    { name: string; url: string | null }
+  >();
 
   // Fetch documents for each data source in parallel.
   await concurrentExecutor(
-    Object.entries(generationsByDataSources),
-    async ([, generations]) => {
-      const dataSource = generations[0].dataSource;
-      const documentIds = [...new Set(generations.map((g) => g.documentId))];
+    Object.entries(docsToFetchByDataSourceId),
+    async ([dataSourceId, documentIds]) => {
+      const dataSource = dataSourceById[dataSourceId];
 
       const docsResult = await coreAPI.getDataSourceDocuments({
         projectId: dataSource.dustAPIProjectId,
@@ -156,7 +173,7 @@ const _sendTrackerWithGenerationEmail = async ({
       }
 
       docsResult.value.documents.forEach((doc) => {
-        documentsById.set(doc.document_id, {
+        documentsByIdentifier.set(`${dataSource.id}__${doc.document_id}`, {
           name: doc.title ?? "Unknown document",
           url: doc.source_url ?? null,
         });
@@ -165,31 +182,59 @@ const _sendTrackerWithGenerationEmail = async ({
     { concurrency: 5 }
   );
 
-  const generationBody = generations.map((generation) => {
-    const doc = documentsById.get(generation.documentId) ?? {
-      name: "Unknown document",
-      url: null,
-    };
-
-    const title = doc.url
-      ? `<a href="${doc.url}" target="_blank">${doc.name}</a>`
-      : `[${doc.name}]`;
-
-    return [
-      `<strong>Changes in document ${title} from ${generation.dataSource.name}:</strong>`,
-      generation.thinking && `<p${generation.thinking}</p>`,
-      `<p>${generation.content}.</p>`,
-    ]
-      .filter(Boolean)
-      .join("");
-  });
+  const generationBody = await Promise.all(
+    generations.map((g) => {
+      // const getDoc = (dataSourceId: number, documentId: string) =>
+      //   documentsByIdentifier.get(`${dataSourceId}__${documentId}`) ?? {
+      //     name: "Unknown document",
+      //     url: null,
+      //   };
+
+      const doc = documentsByIdentifier.get(
+        `${g.dataSource.id}__${g.documentId}`
+      ) ?? {
+        name: "Unknown document",
+        url: null,
+      };
+      const maintainedDoc = g.maintainedDataSource
+        ? documentsByIdentifier.get(
+            `${g.maintainedDataSource.id}__${g.maintainedDocumentId}`
+          ) ?? null
+        : null;
+
+      const title = doc.url
+        ? `<a href="${doc.url}" target="_blank">${doc.name}</a>`
+        : `[${doc.name}]`;
+
+      let maintainedTitle: string | null = null;
+      if (maintainedDoc) {
+        maintainedTitle = maintainedDoc.url
+          ? `<a href="${maintainedDoc.url}" target="_blank">${maintainedDoc.name}</a>`
+          : `[${maintainedDoc.name}]`;
+      }
+
+      let body = `<strong>Changes in document ${title} from ${g.dataSource.name}`;
+      if (maintainedTitle && g.maintainedDataSource) {
+        body += ` might affect ${maintainedTitle} from ${g.maintainedDataSource.name}`;
+      }
+      body += `:</strong>`;
+
+      // TODO(@fontanierh): maybe add thinking back in.
+      // if (g.thinking) {
+      //   body += `<p${g.thinking}</p>`;
+      // }
+
+      body += `<p>${g.content.replace(/\n/g, "<br />")}.</p>`;
+      return body;
+    })
+  );
 
   const body = `
 <p>We have new suggestions for your tracker ${name}:</p>
 <p>${generations.length} recommendations were generated due to changes in watched documents.</p>
 <br />
 <br />
-${generationBody.join("<br />")}
+${generationBody.join("<hr />")}
 `;
 
   await sendEmailWithTemplate({

diff --git a/front/lib/models/doc_tracker.ts b/front/lib/models/doc_tracker.ts
@@ -237,14 +237,16 @@ export class TrackerGenerationModel extends SoftDeletableModel<TrackerGeneration
   declare trackerConfigurationId: ForeignKey<TrackerConfigurationModel["id"]>;
   declare dataSourceId: ForeignKey<DataSourceModel["id"]>;
   declare documentId: string;
-  declare maintainedDocumentDataSourceId: ForeignKey<DataSourceModel["id"]>;
-  declare maintainedDocumentId: string;
+  declare maintainedDocumentDataSourceId: ForeignKey<
+    DataSourceModel["id"]
+  > | null;
+  declare maintainedDocumentId: string | null;
 
   declare consumedAt: Date | null;
 
   declare trackerConfiguration: NonAttribute<TrackerConfigurationModel>;
   declare dataSource: NonAttribute<DataSourceModel>;
-  declare maintainedDocumentDataSource: NonAttribute<DataSourceModel>;
+  declare maintainedDocumentDataSource: NonAttribute<DataSourceModel> | null;
 }
 
 TrackerGenerationModel.init(

diff --git a/front/lib/resources/tracker_resource.ts b/front/lib/resources/tracker_resource.ts
@@ -437,6 +437,11 @@ export class TrackerConfigurationResource extends ResourceWithSpace<TrackerConfi
               as: "dataSource",
               required: true,
             },
+            {
+              model: DataSourceModel,
+              as: "maintainedDocumentDataSource",
+              required: false,
+            },
           ],
         },
       ],
@@ -701,6 +706,17 @@ export class TrackerConfigurationResource extends ResourceWithSpace<TrackerConfi
             dustAPIProjectId: g.dataSource.dustAPIProjectId,
             dustAPIDataSourceId: g.dataSource.dustAPIDataSourceId,
           },
+          maintainedDataSource: g.maintainedDocumentDataSource
+            ? {
+                id: g.maintainedDocumentDataSource.id,
+                name: dataSourceName,
+                dustAPIProjectId:
+                  g.maintainedDocumentDataSource.dustAPIProjectId,
+                dustAPIDataSourceId:
+                  g.maintainedDocumentDataSource.dustAPIDataSourceId,
+              }
+            : null,
+          maintainedDocumentId: g.maintainedDocumentId,
         };
       });
     }

diff --git a/types/src/front/tracker.ts b/types/src/front/tracker.ts
@@ -68,15 +68,19 @@ export type TrackerIdWorkspaceId = {
   workspaceId: string;
 };
 
+export type TrackerDataSource = {
+  id: ModelId;
+  name: string;
+  dustAPIProjectId: string;
+  dustAPIDataSourceId: string;
+};
+
 export type TrackerGenerationToProcess = {
   id: ModelId;
   content: string;
   thinking: string | null;
   documentId: string;
-  dataSource: {
-    id: ModelId;
-    name: string;
-    dustAPIProjectId: string;
-    dustAPIDataSourceId: string;
-  };
+  dataSource: TrackerDataSource;
+  maintainedDataSource: TrackerDataSource | null;
+  maintainedDocumentId: string | null;
 };