Replaying DataSourceCreated and AssistantCreated

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";

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;
 }

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<Authenticator> {
+    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<RoleType> => {
+          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.

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);