Use WorkspaceAwareModel on conversation related models (#10184)

* Use WorkspaceAwareModel on conversation related models * Add migrations * Fix migration * 📖 * Backfill content fragments * 👕 * ➡️
dust-tt · Jan 23, 2025 · 73e3af2 · 73e3af2
1 parent 92fd352
commit 73e3af2
Show file tree

Hide file tree

Showing 8 changed files with 301 additions and 19 deletions.
diff --git a/front/lib/api/assistant/conversation.ts b/front/lib/api/assistant/conversation.ts
@@ -316,6 +316,7 @@ async function createOrUpdateParticipation({
         conversationId: conversation.id,
         action: "posted",
         userId: user.id,
+        workspaceId: conversation.owner.id,
       });
     }
   }
@@ -774,7 +775,7 @@ export async function* postUserMessage(
           transaction: t,
         })) ?? -1) + 1;
 
-      async function createMessageAndUserMessage() {
+      async function createMessageAndUserMessage(workspace: WorkspaceType) {
         return Message.create(
           {
             sId: generateRandomModelSId(),
@@ -795,22 +796,24 @@ export async function* postUserMessage(
                     ? user.id
                     : (
                         await attributeUserFromWorkspaceAndEmail(
-                          owner,
+                          workspace,
                           context.email
                         )
                       )?.id,
+                  workspaceId: workspace.id,
                 },
                 { transaction: t }
               )
             ).id,
+            workspaceId: workspace.id,
           },
           {
             transaction: t,
           }
         );
       }
 
-      const m = await createMessageAndUserMessage();
+      const m = await createMessageAndUserMessage(owner);
       const userMessage: UserMessageWithRankType = {
         id: m.id,
         created: m.createdAt.getTime(),
@@ -843,6 +846,7 @@ export async function* postUserMessage(
                 {
                   messageId: m.id,
                   agentConfigurationId: configuration.sId,
+                  workspaceId: owner.id,
                 },
                 { transaction: t }
               );
@@ -852,6 +856,7 @@ export async function* postUserMessage(
                   status: "created",
                   agentConfigurationId: configuration.sId,
                   agentConfigurationVersion: configuration.version,
+                  workspaceId: owner.id,
                 },
                 { transaction: t }
               );
@@ -862,6 +867,7 @@ export async function* postUserMessage(
                   conversationId: conversation.id,
                   parentId: userMessage.id,
                   agentMessageId: agentMessageRow.id,
+                  workspaceId: owner.id,
                 },
                 {
                   transaction: t,
@@ -1236,7 +1242,10 @@ export async function* editUserMessage(
       }
       const userMessageRow = messageRow.userMessage;
       // adding messageRow as param otherwise Ts doesn't get it can't be null
-      async function createMessageAndUserMessage(messageRow: Message) {
+      async function createMessageAndUserMessage(
+        workspace: WorkspaceType,
+        messageRow: Message
+      ) {
         return Message.create(
           {
             sId: generateRandomModelSId(),
@@ -1259,22 +1268,24 @@ export async function* editUserMessage(
                     ? userMessageRow.userId
                     : (
                         await attributeUserFromWorkspaceAndEmail(
-                          owner,
+                          workspace,
                           userMessageRow.userContextEmail
                         )
                       )?.id,
+                  workspaceId: workspace.id,
                 },
                 { transaction: t }
               )
             ).id,
+            workspaceId: workspace.id,
           },
           {
             transaction: t,
           }
         );
       }
 
-      const m = await createMessageAndUserMessage(messageRow);
+      const m = await createMessageAndUserMessage(owner, messageRow);
 
       const userMessage: UserMessageWithRankType = {
         id: m.id,
@@ -1320,6 +1331,7 @@ export async function* editUserMessage(
               {
                 messageId: m.id,
                 agentConfigurationId: configuration.sId,
+                workspaceId: owner.id,
               },
               { transaction: t }
             );
@@ -1329,6 +1341,7 @@ export async function* editUserMessage(
                 status: "created",
                 agentConfigurationId: configuration.sId,
                 agentConfigurationVersion: configuration.version,
+                workspaceId: owner.id,
               },
               { transaction: t }
             );
@@ -1339,6 +1352,7 @@ export async function* editUserMessage(
                 conversationId: conversation.id,
                 parentId: userMessage.id,
                 agentMessageId: agentMessageRow.id,
+                workspaceId: owner.id,
               },
               {
                 transaction: t,
@@ -1568,6 +1582,7 @@ export async function* retryAgentMessage(
           agentConfigurationId: messageRow.agentMessage.agentConfigurationId,
           agentConfigurationVersion:
             messageRow.agentMessage.agentConfigurationVersion,
+          workspaceId: auth.getNonNullableWorkspace().id,
         },
         { transaction: t }
       );
@@ -1579,6 +1594,7 @@ export async function* retryAgentMessage(
           parentId: messageRow.parentId,
           version: messageRow.version + 1,
           agentMessageId: agentMessageRow.id,
+          workspaceId: auth.getNonNullableWorkspace().id,
         },
         {
           transaction: t,
@@ -1759,6 +1775,7 @@ export async function postNewContentFragment(
         userContextEmail: context?.email,
         userContextFullName: context?.fullName,
         userContextUsername: context?.username,
+        workspaceId: owner.id,
       };
 
       const contentFragment = await (() => {
@@ -1786,6 +1803,7 @@ export async function postNewContentFragment(
           rank: nextMessageRank,
           conversationId: conversation.id,
           contentFragmentId: contentFragment.id,
+          workspaceId: owner.id,
         },
         {
           transaction: t,

diff --git a/front/lib/api/assistant/reaction.ts b/front/lib/api/assistant/reaction.ts
@@ -106,10 +106,7 @@ export async function createMessageReaction(
     reaction: string;
   }
 ): Promise<boolean | null> {
-  const owner = auth.workspace();
-  if (!owner) {
-    throw new Error("Unexpected `auth` without `workspace`.");
-  }
+  const owner = auth.getNonNullableWorkspace();
 
   const message = await Message.findOne({
     where: {
@@ -128,6 +125,7 @@ export async function createMessageReaction(
     userContextUsername: context.username,
     userContextFullName: context.fullName,
     reaction,
+    workspaceId: owner.id,
   });
   return newReaction !== null;
 }

diff --git a/front/lib/models/assistant/conversation.ts b/front/lib/models/assistant/conversation.ts
@@ -13,7 +13,6 @@ import type { AgentMessageContent } from "@app/lib/models/assistant/agent_messag
 import { frontSequelize } from "@app/lib/resources/storage";
 import { ContentFragmentModel } from "@app/lib/resources/storage/models/content_fragment";
 import { UserModel } from "@app/lib/resources/storage/models/user";
-import { BaseModel } from "@app/lib/resources/storage/wrappers/base";
 import { WorkspaceAwareModel } from "@app/lib/resources/storage/wrappers/workspace_models";
 
 export class Conversation extends WorkspaceAwareModel<Conversation> {
@@ -84,7 +83,7 @@ Conversation.init(
   }
 );
 
-export class ConversationParticipant extends BaseModel<ConversationParticipant> {
+export class ConversationParticipant extends WorkspaceAwareModel<ConversationParticipant> {
   declare createdAt: CreationOptional<Date>;
   declare updatedAt: CreationOptional<Date>;
 
@@ -150,7 +149,7 @@ ConversationParticipant.belongsTo(UserModel, {
   foreignKey: { name: "userId", allowNull: false },
 });
 
-export class UserMessage extends BaseModel<UserMessage> {
+export class UserMessage extends WorkspaceAwareModel<UserMessage> {
   declare createdAt: CreationOptional<Date>;
   declare updatedAt: CreationOptional<Date>;
 
@@ -220,7 +219,7 @@ UserMessage.belongsTo(UserModel, {
   foreignKey: { name: "userId", allowNull: true },
 });
 
-export class AgentMessage extends BaseModel<AgentMessage> {
+export class AgentMessage extends WorkspaceAwareModel<AgentMessage> {
   declare createdAt: CreationOptional<Date>;
   declare updatedAt: CreationOptional<Date>;
   declare runIds: string[] | null;
@@ -371,7 +370,7 @@ AgentMessageFeedback.belongsTo(AgentMessage, {
   as: "agentMessage",
 });
 
-export class Message extends BaseModel<Message> {
+export class Message extends WorkspaceAwareModel<Message> {
   declare createdAt: CreationOptional<Date>;
   declare updatedAt: CreationOptional<Date>;
 
@@ -512,7 +511,7 @@ Message.belongsTo(ContentFragmentModel, {
   foreignKey: { name: "contentFragmentId", allowNull: true },
 });
 
-export class MessageReaction extends BaseModel<MessageReaction> {
+export class MessageReaction extends WorkspaceAwareModel<MessageReaction> {
   declare createdAt: CreationOptional<Date>;
   declare updatedAt: CreationOptional<Date>;
 
@@ -583,7 +582,7 @@ MessageReaction.belongsTo(UserModel, {
   foreignKey: { name: "userId", allowNull: true }, // null = mention is not a user using a Slackbot
 });
 
-export class Mention extends BaseModel<Mention> {
+export class Mention extends WorkspaceAwareModel<Mention> {
   declare createdAt: CreationOptional<Date>;
   declare updatedAt: CreationOptional<Date>;
 

diff --git a/front/lib/resources/content_fragment_resource.ts b/front/lib/resources/content_fragment_resource.ts
@@ -59,6 +59,7 @@ export class ContentFragmentResource extends BaseResource<ContentFragmentModel>
         ...blob,
         sId: generateRandomModelSId("cf"),
         version: "latest",
+        workspaceId: blob.workspaceId,
       },
       {
         transaction,
@@ -91,6 +92,7 @@ export class ContentFragmentResource extends BaseResource<ContentFragmentModel>
           ...blob,
           sId,
           version: "latest",
+          workspaceId: blob.workspaceId,
         },
         {
           transaction: t,

diff --git a/front/lib/resources/storage/models/content_fragment.ts b/front/lib/resources/storage/models/content_fragment.ts
@@ -8,9 +8,9 @@ import { DataTypes } from "sequelize";
 import { frontSequelize } from "@app/lib/resources/storage";
 import { FileModel } from "@app/lib/resources/storage/models/files";
 import { UserModel } from "@app/lib/resources/storage/models/user";
-import { BaseModel } from "@app/lib/resources/storage/wrappers/base";
+import { WorkspaceAwareModel } from "@app/lib/resources/storage/wrappers/workspace_models";
 
-export class ContentFragmentModel extends BaseModel<ContentFragmentModel> {
+export class ContentFragmentModel extends WorkspaceAwareModel<ContentFragmentModel> {
   declare createdAt: CreationOptional<Date>;
   declare updatedAt: CreationOptional<Date>;