Skip to content

Commit

Permalink
scrub/delete workflows: share conversation deletion logic (#9643)
Browse files Browse the repository at this point in the history
* scrub/delete workflows: share conversation deletion logic

* nit
  • Loading branch information
spolu authored Dec 30, 2024
1 parent 337094a commit a90b240
Show file tree
Hide file tree
Showing 3 changed files with 9 additions and 150 deletions.
8 changes: 4 additions & 4 deletions front/lib/api/assistant/conversation/destroy.ts
Original file line number Diff line number Diff line change
Expand Up @@ -127,14 +127,14 @@ async function destroyContentFragments(
}

async function destroyConversationDataSource({
workspaceId,
workspace,
conversationId,
}: {
workspaceId: string;
workspace: LightWorkspaceType;
conversationId: string;
}) {
// We need an authenticator to interact with the data source resource.
const auth = await Authenticator.internalAdminForWorkspace(workspaceId);
const auth = await Authenticator.internalAdminForWorkspace(workspace.sId);
const conversation = await getConversationWithoutContent(
auth,
conversationId,
Expand Down Expand Up @@ -216,7 +216,7 @@ export async function destroyConversation(
});

await destroyConversationDataSource({
workspaceId: workspace.sId,
workspace,
conversationId: conversation.sId,
});

Expand Down
144 changes: 3 additions & 141 deletions front/poke/temporal/activities.ts
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ import { TrackerConfigurationResource } from "@app/lib/resources/tracker_resourc
import { UserResource } from "@app/lib/resources/user_resource";
import { renderLightWorkspaceType } from "@app/lib/workspace";
import logger from "@app/logger/logger";
import { deleteAllConversations } from "@app/temporal/scrub_workspace/activities";

const hardDeleteLogger = logger.child({ activity: "hard-delete" });

Expand Down Expand Up @@ -172,7 +173,7 @@ export async function isWorkflowDeletableActivity({
workspaceId: string;
}) {
const auth = await Authenticator.internalAdminForWorkspace(workspaceId);
const workspace = await auth.getNonNullableWorkspace();
const workspace = auth.getNonNullableWorkspace();

return areAllSubscriptionsCanceled(renderLightWorkspaceType({ workspace }));
}
Expand All @@ -183,146 +184,7 @@ export async function deleteConversationsActivity({
workspaceId: string;
}) {
const auth = await Authenticator.internalAdminForWorkspace(workspaceId);
const workspace = auth.workspace();

if (!workspace) {
throw new Error("Could not find the workspace.");
}

const conversations = await Conversation.findAll({
where: {
workspaceId: workspace.id,
},
});
const chunkSize = 8;
const chunks: Conversation[][] = [];
for (let i = 0; i < conversations.length; i += chunkSize) {
chunks.push(conversations.slice(i, i + chunkSize));
}

await frontSequelize.transaction(async (t) => {
for (let i = 0; i < chunks.length; i++) {
const chunk = chunks[i];
if (!chunk) {
continue;
}
await Promise.all(
chunk.map((c) => {
return (async (): Promise<void> => {
const messages = await Message.findAll({
where: { conversationId: c.id },
transaction: t,
});
for (const msg of messages) {
if (msg.userMessageId) {
await UserMessage.destroy({
where: { id: msg.userMessageId },
transaction: t,
});
}
if (msg.agentMessageId) {
const agentMessage = await AgentMessage.findOne({
where: { id: msg.agentMessageId },
transaction: t,
});
if (agentMessage) {
const retrievalAction = await AgentRetrievalAction.findOne({
where: {
agentMessageId: agentMessage.id,
},
transaction: t,
});
if (retrievalAction) {
await RetrievalDocumentResource.deleteAllForActions([
retrievalAction.id,
]);

await AgentRetrievalAction.destroy({
where: { id: retrievalAction.id },
transaction: t,
});
}

await AgentMessageContent.destroy({
where: { agentMessageId: agentMessage.id },
transaction: t,
});

await AgentMessageFeedback.destroy({
where: { agentMessageId: agentMessage.id },
transaction: t,
});

// Delete associated actions.

await AgentBrowseAction.destroy({
where: { agentMessageId: agentMessage.id },
transaction: t,
});

await AgentProcessAction.destroy({
where: { agentMessageId: agentMessage.id },
transaction: t,
});

await AgentTablesQueryAction.destroy({
where: { agentMessageId: agentMessage.id },
transaction: t,
});

await AgentWebsearchAction.destroy({
where: { agentMessageId: agentMessage.id },
transaction: t,
});

await agentMessage.destroy({ transaction: t });
}
}
if (msg.contentFragmentId) {
const contentFragment =
await ContentFragmentResource.fetchByModelId(
msg.contentFragmentId,
t
);
if (contentFragment) {
await contentFragment.destroy(
{
conversationId: c.sId,
messageId: msg.sId,
workspaceId: workspace.sId,
},
t
);
}
}
await MessageReaction.destroy({
where: { messageId: msg.id },
transaction: t,
});
await Mention.destroy({
where: { messageId: msg.id },
transaction: t,
});
await msg.destroy({ transaction: t });
}
await ConversationParticipant.destroy({
where: { conversationId: c.id },
transaction: t,
});

hardDeleteLogger.info(
{
conversationId: c.sId,
},
"Deleting conversation"
);

await c.destroy({ transaction: t });
})();
})
);
}
});
await deleteAllConversations(auth);
}

export async function deleteAgentsActivity({
Expand Down
7 changes: 2 additions & 5 deletions front/temporal/scrub_workspace/activities.ts
Original file line number Diff line number Diff line change
Expand Up @@ -115,11 +115,8 @@ export async function pauseAllConnectors({
}
}

async function deleteAllConversations(auth: Authenticator) {
const workspace = auth.workspace();
if (!workspace) {
throw new Error("No workspace found");
}
export async function deleteAllConversations(auth: Authenticator) {
const workspace = auth.getNonNullableWorkspace();
const conversations = await Conversation.findAll({
where: { workspaceId: workspace.id },
});
Expand Down

0 comments on commit a90b240

Please sign in to comment.