diff --git a/front/lib/api/tracker.ts b/front/lib/api/tracker.ts index 38627bf4f6c6..bddcb86984f9 100644 --- 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 => { 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 - ? `${doc.name}` - : `[${doc.name}]`; - - return [ - `Changes in document ${title} from ${generation.dataSource.name}:`, - generation.thinking && ``, - `

${generation.content}.

`, - ] - .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 + ? `${doc.name}` + : `[${doc.name}]`; + + let maintainedTitle: string | null = null; + if (maintainedDoc) { + maintainedTitle = maintainedDoc.url + ? `${maintainedDoc.name}` + : `[${maintainedDoc.name}]`; + } + + let body = `Changes in document ${title} from ${g.dataSource.name}`; + if (maintainedTitle && g.maintainedDataSource) { + body += ` might affect ${maintainedTitle} from ${g.maintainedDataSource.name}`; + } + body += `:`; + + if (g.thinking) { + body += ` +
+ View thinking +

${g.thinking.replace(/\n/g, "
")}

+
`; + } + + body += `

${g.content.replace(/\n/g, "
")}.

`; + return body; + }) + ); const body = `

We have new suggestions for your tracker ${name}:

${generations.length} recommendations were generated due to changes in watched documents.



-${generationBody.join("
")} +${generationBody.join("
")} `; await sendEmailWithTemplate({ diff --git a/front/lib/models/doc_tracker.ts b/front/lib/models/doc_tracker.ts index 3b9af1c38b21..5e40b00c0265 100644 --- a/front/lib/models/doc_tracker.ts +++ b/front/lib/models/doc_tracker.ts @@ -237,14 +237,16 @@ export class TrackerGenerationModel extends SoftDeletableModel; declare dataSourceId: ForeignKey; declare documentId: string; - declare maintainedDocumentDataSourceId: ForeignKey; - declare maintainedDocumentId: string; + declare maintainedDocumentDataSourceId: ForeignKey< + DataSourceModel["id"] + > | null; + declare maintainedDocumentId: string | null; declare consumedAt: Date | null; declare trackerConfiguration: NonAttribute; declare dataSource: NonAttribute; - declare maintainedDocumentDataSource: NonAttribute; + declare maintainedDocumentDataSource: NonAttribute | null; } TrackerGenerationModel.init( diff --git a/front/lib/resources/tracker_resource.ts b/front/lib/resources/tracker_resource.ts index 7910e389db46..3a197b57c6c9 100644 --- a/front/lib/resources/tracker_resource.ts +++ b/front/lib/resources/tracker_resource.ts @@ -437,6 +437,11 @@ export class TrackerConfigurationResource extends ResourceWithSpace { + 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 index 9eb0e275b0a1..74f960c4e6fe 100644 --- 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; };