feat: add link to maintained document in tracker email (#9944)

Co-authored-by: Henry Fontanier <[email protected]>
dust-tt · Jan 14, 2025 · d857ad6 · d857ad6
1 parent 3068af8
commit d857ad6
Show file tree

Hide file tree

Showing 5 changed files with 205 additions and 40 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";
@@ -84,8 +84,8 @@ const sendTrackerEmail = async ({
 
   const sendEmail =
     generations.length > 0
-      ? _sendTrackerWithGenerationEmail
-      : _sendTrackerDefaultEmail;
+      ? sendTrackerWithGenerationEmail
+      : sendTrackerDefaultEmail;
 
   await Promise.all(
     Array.from(recipients).map((recipient) =>
@@ -94,7 +94,7 @@ const sendTrackerEmail = async ({
   );
 };
 
-const _sendTrackerDefaultEmail = async ({
+const sendTrackerDefaultEmail = async ({
   name,
   recipient,
 }: {
@@ -115,7 +115,7 @@ const _sendTrackerDefaultEmail = async ({
   });
 };
 
-const _sendTrackerWithGenerationEmail = async ({
+export const sendTrackerWithGenerationEmail = async ({
   name,
   recipient,
   generations,
@@ -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,56 @@ 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 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>`;
+
+      if (g.thinking) {
+        body += `
+          <details>
+            <summary>View thinking</summary>
+            <p>${g.thinking.replace(/\n/g, "<br />")}</p>
+          </details>`;
+      }
+
+      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/front/scripts/send_tracker_generations.ts b/front/scripts/send_tracker_generations.ts
@@ -0,0 +1,101 @@
+import { sendTrackerWithGenerationEmail } from "@app/lib/api/tracker";
+import { TrackerGenerationModel } from "@app/lib/models/doc_tracker";
+import { frontSequelize } from "@app/lib/resources/storage";
+import { DataSourceModel } from "@app/lib/resources/storage/models/data_source";
+import { isEmailValid } from "@app/lib/utils";
+import { makeScript } from "@app/scripts/helpers";
+
+makeScript(
+  {
+    generationIds: {
+      type: "array",
+      demandOption: true,
+      description: "List of generation IDs",
+    },
+    email: {
+      type: "string",
+      demandOption: true,
+      description: "Email address to send to",
+    },
+  },
+  async ({ execute, generationIds, email }, logger) => {
+    try {
+      // Validate email
+      if (!isEmailValid(email)) {
+        throw new Error("Invalid email address");
+      }
+
+      // Parse and validate generation IDs
+      const ids = generationIds.map((id) => parseInt(id));
+      if (ids.some((id) => isNaN(id))) {
+        throw new Error("Invalid generation IDs - must be numbers");
+      }
+
+      if (execute) {
+        // Fetch generations with their data sources
+        const generations = await TrackerGenerationModel.findAll({
+          where: {
+            id: ids,
+          },
+          include: [
+            {
+              model: DataSourceModel,
+              required: true,
+            },
+            {
+              model: DataSourceModel,
+              as: "maintainedDocumentDataSource",
+              required: false,
+            },
+          ],
+        });
+
+        if (generations.length === 0) {
+          throw new Error("No generations found with the provided IDs");
+        }
+
+        // Convert to TrackerGenerationToProcess format
+        const generationsToProcess = generations.map((g) => ({
+          id: g.id,
+          content: g.content,
+          thinking: g.thinking,
+          documentId: g.documentId,
+          dataSource: {
+            id: g.dataSource.id,
+            name: g.dataSource.name,
+            dustAPIProjectId: g.dataSource.dustAPIProjectId,
+            dustAPIDataSourceId: g.dataSource.dustAPIDataSourceId,
+          },
+          maintainedDocumentId: g.maintainedDocumentId,
+          maintainedDataSource: g.maintainedDocumentDataSource
+            ? {
+                id: g.maintainedDocumentDataSource.id,
+                name: g.maintainedDocumentDataSource.name,
+                dustAPIProjectId:
+                  g.maintainedDocumentDataSource.dustAPIProjectId,
+                dustAPIDataSourceId:
+                  g.maintainedDocumentDataSource.dustAPIDataSourceId,
+              }
+            : null,
+        }));
+
+        // Send email
+        await sendTrackerWithGenerationEmail({
+          name: "Manual Generation Email",
+          recipient: email,
+          generations: generationsToProcess,
+          localLogger: logger,
+        });
+
+        logger.info({}, "Email sent successfully");
+      } else {
+        logger.info(
+          { generationIds: ids, email },
+          "Dry run - would send email with these parameters"
+        );
+      }
+    } finally {
+      await frontSequelize.close();
+    }
+  }
+);
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;
 };