diff --git a/front/documents_post_process_hooks/hooks/document_tracker/suggest_changes/lib.ts b/front/documents_post_process_hooks/hooks/document_tracker/suggest_changes/lib.ts index 8f4b7460e4df8..a57140384106f 100644 --- a/front/documents_post_process_hooks/hooks/document_tracker/suggest_changes/lib.ts +++ b/front/documents_post_process_hooks/hooks/document_tracker/suggest_changes/lib.ts @@ -3,7 +3,7 @@ import sgMail from "@sendgrid/mail"; import { Op } from "sequelize"; import showdown from "showdown"; -import { +import type { DocumentsPostProcessHookFilterParams, DocumentsPostProcessHookOnUpsertParams, } from "@app/documents_post_process_hooks/hooks"; diff --git a/front/lib/amplitude/back/index.ts b/front/lib/amplitude/back/index.ts index 26395d768efd1..e665b1a794f47 100644 --- a/front/lib/amplitude/back/index.ts +++ b/front/lib/amplitude/back/index.ts @@ -144,7 +144,7 @@ export function trackUserMessage({ ); } -export function trackDataSourceCreated( +export async function trackDataSourceCreated( auth: Authenticator, { dataSource, @@ -166,17 +166,18 @@ export function trackDataSourceCreated( assistantDefaultSelected: dataSource.assistantDefaultSelected, }); - amplitude.track( + await amplitude.track( `user-${userId}`, { ...event, groups: { [GROUP_TYPE]: workspace.sId } }, { time: dataSource.createdAt, insert_id: `data_source_created_${dataSource.id}`, } - ); + ).promise; + console.log("tracking data source created", dataSource.name, workspace.sId); } -export function trackAssistantCreated( +export async function trackAssistantCreated( auth: Authenticator, { assistant }: { assistant: AgentConfigurationType } ) { @@ -195,7 +196,7 @@ export function trackAssistantCreated( assistantActionType: assistant.action?.type || "", assistantVersion: assistant.version, }); - amplitude.track( + await amplitude.track( `user-${userId}`, { ...event, groups: { [GROUP_TYPE]: workspace.sId } }, { @@ -204,5 +205,5 @@ export function trackAssistantCreated( : Date.now(), insert_id: `assistant_created_${assistant.sId}`, } - ); + ).promise; } diff --git a/front/lib/auth.ts b/front/lib/auth.ts index 28974ca26eb6a..dc9760aba4b1f 100644 --- a/front/lib/auth.ts +++ b/front/lib/auth.ts @@ -1,5 +1,6 @@ import { getSession as getAuth0Session } from "@auth0/nextjs-auth0"; import type { + ModelId, RoleType, UserType, WhitelistableFeature, @@ -160,6 +161,64 @@ export class Authenticator { }); } + static async fromIds(userId: ModelId, wId: string): Promise { + const [workspace, user] = await Promise.all([ + (async () => { + return Workspace.findOne({ + where: { + sId: wId, + }, + }); + })(), + (async () => { + return User.findOne({ + where: { + id: userId, + }, + }); + })(), + ]); + + let role = "none" as RoleType; + let subscription: SubscriptionType | null = null; + let flags: WhitelistableFeature[] = []; + + if (user && workspace) { + [role, subscription, flags] = await Promise.all([ + (async (): Promise => { + const membership = await Membership.findOne({ + where: { + userId: user.id, + workspaceId: workspace.id, + }, + }); + return membership && + ["admin", "builder", "user"].includes(membership.role) + ? (membership.role as RoleType) + : "none"; + })(), + subscriptionForWorkspace(workspace.sId), + (async () => { + return ( + await FeatureFlag.findAll({ + where: { + workspaceId: workspace.id, + }, + }) + ).map((flag) => flag.name); + })(), + ]); + } + + return new Authenticator({ + workspace, + user, + role, + subscription, + flags, + }); + } + /** * Get a an Authenticator for the target workspace and the authentified Super User user from the * Auth0 session. diff --git a/front/migrations/20240325_amplitude_replay_assistants_datasources.ts b/front/migrations/20240325_amplitude_replay_assistants_datasources.ts new file mode 100644 index 0000000000000..6b09d0a09c987 --- /dev/null +++ b/front/migrations/20240325_amplitude_replay_assistants_datasources.ts @@ -0,0 +1,98 @@ +import type { AgentStatus } from "@dust-tt/types"; + +import { + populateWorkspaceProperties, + trackAssistantCreated, + trackDataSourceCreated, +} from "@app/lib/amplitude/back"; +import { getAgentConfiguration } from "@app/lib/api/assistant/configuration"; +import { getDataSource } from "@app/lib/api/data_sources"; +import { Authenticator } from "@app/lib/auth"; +import { + AgentConfiguration, + Conversation, + DataSource, + Membership, + Workspace, +} from "@app/lib/models"; + +const workspaceIDS = ["c2d481e559"]; + +async function populateAssistantCreated(workspace: Workspace) { + const assistants = await AgentConfiguration.findAll({ + where: { + workspaceId: workspace.id, + status: ["active", "archived"] satisfies AgentStatus[], + }, + }); + for (const assistant of assistants) { + const auth = await Authenticator.fromIds(assistant.authorId, workspace.sId); + if (!auth.isUser()) { + throw new Error("Only users can create agents."); + } + if (!auth.workspace()) { + throw new Error("Workspace not found."); + } + const agentConfigType = await getAgentConfiguration(auth, assistant.sId); + if (!agentConfigType) { + throw new Error("Agent configuration not found." + assistant.sId); + } + await trackAssistantCreated(auth, { assistant: agentConfigType }); + console.log("tracked assistant created", assistant.sId, assistant.name); + } +} + +export async function populateDataSourceCreated(workspace: Workspace) { + const dataSources = await DataSource.findAll({ + where: { + workspaceId: workspace.id, + }, + }); + const defaultAdmin = await Membership.findOne({ + where: { + workspaceId: workspace.id, + role: "admin", + }, + order: [["createdAt", "ASC"]], + limit: 1, + }); + for (const dataSource of dataSources) { + const adminId = dataSource.editedByUserId || defaultAdmin?.userId; + if (!adminId) { + throw new Error("Admin not found."); + } + const auth = await Authenticator.fromIds(adminId, workspace.sId); + if (!auth.isUser()) { + throw new Error("Only users can create agents."); + } + if (!auth.workspace()) { + throw new Error("Workspace not found."); + } + const ds = await getDataSource(auth, dataSource.name); + if (!ds) { + throw new Error( + `Data source not found: ${dataSource.name} in workspace ${workspace.sId}` + ); + } + await trackDataSourceCreated(auth, { + dataSource: ds, + }); + } +} + +async function main() { + for (const workspaceID of workspaceIDS) { + const workspace = await Workspace.findOne({ + where: { + sId: workspaceID, + }, + }); + if (!workspace) { + throw new Error(`Workspace not found: ${workspaceID}`); + } + await populateAssistantCreated(workspace); + await populateDataSourceCreated(workspace); + } +} + +main().catch(console.error);