Classify conversations instead of user_messages

front/admin/db.ts

@@ -45,9 +45,9 @@ import {
   GlobalAgentSettings,
 } from "@app/lib/models/assistant/agent";
 import { ContentFragment } from "@app/lib/models/assistant/conversation";
+import { ConversationClassification } from "@app/lib/models/conversation_classification";
 import { FeatureFlag } from "@app/lib/models/feature_flag";
 import { PlanInvitation } from "@app/lib/models/plan";
-import { UserMessageClassification } from "@app/lib/models/user_message_classification";
 
 async function main() {
   await User.sync({ alter: true });
@@ -100,7 +100,7 @@ async function main() {
 
   await FeatureFlag.sync({ alter: true });
 
-  await UserMessageClassification.sync({ alter: true });
+  await ConversationClassification.sync({ alter: true });
   process.exit(0);
 }
 

front/admin/tools/message_classification.ts

@@ -3,9 +3,9 @@ import OpenAI from "openai";
 
 import { UserMessage } from "@app/lib/models";
 import { Conversation, Message, Workspace } from "@app/lib/models";
-import { UserMessageClassification } from "@app/lib/models/user_message_classification";
+import { ConversationClassification } from "@app/lib/models/conversation_classification";
 
-async function classifyUserMessage(userMessage: UserMessage) {
+async function classifyConversation(content: string) {
   if (!process.env.DUST_MANAGED_OPENAI_API_KEY) {
     throw new Error("DUST_MANAGED_OPENAI_API_KEY is not set");
   }
@@ -16,7 +16,7 @@ async function classifyUserMessage(userMessage: UserMessage) {
   const prompt = `Classify this message as one class of the following classes: ${MESSAGE_CLASSES.join(
     ", "
   )}:`;
-  const promptWithMessage = `${prompt}\n${userMessage.content}`;
+  const promptWithMessage = `${prompt}\n${content}`;
   const chatCompletion = await openai.chat.completions.create({
     messages: [{ role: "user", content: promptWithMessage }],
     model: "gpt-3.5-turbo",
@@ -53,7 +53,6 @@ export async function classifyWorkspace({
   workspaceId: string;
   limit: number;
 }) {
-  let count = 0;
   const workspace = await Workspace.findOne({
     where: {
       sId: workspaceId,
@@ -71,46 +70,63 @@ export async function classifyWorkspace({
     limit: limit,
     order: [["id", "DESC"]],
   });
+  console.log("conversations", conversations.length);
 
   for (const conversation of conversations) {
     const messages = await Message.findAll({
       where: {
         conversationId: conversation.id,
       },
       attributes: ["id", "userMessageId"],
+      order: [["id", "ASC"]],
     });
+    if (messages.length > 30) {
+      console.log("too many messages", conversation.id);
+      continue;
+    }
+    const renderedConversation: { username: string; content: string }[] = [];
     for (const message of messages) {
       if (message.userMessageId) {
         const userMessage = await UserMessage.findByPk(message.userMessageId);
-        if (userMessage) {
-          if (
-            await UserMessageClassification.findOne({
-              where: { userMessageId: userMessage.id },
-            })
-          ) {
-            console.log("already classified", userMessage.id);
-            continue;
-          }
-          const result = await classifyUserMessage(userMessage);
-          console.log(
-            `[%s] [%s]`,
-            userMessage.content.substring(0, 10),
-            result
-          );
-          if (result && isMessageClassification(result)) {
-            await UserMessageClassification.upsert({
-              messageClass: result,
-              userMessageId: userMessage.id,
-            });
-            count++;
-            if (count >= limit) {
-              console.log("limit reached");
-              return;
-            }
-          } else {
-            console.log("could not classify message", userMessage.id);
-          }
+        if (!userMessage) {
+          console.log("user message not found", message.userMessageId);
+          continue;
         }
+        renderedConversation.push({
+          username:
+            userMessage.userContextFullName ||
+            userMessage.userContextEmail ||
+            userMessage.userContextUsername,
+          content: userMessage.content,
+        });
+      }
+    }
+    if (renderedConversation.length > 0) {
+      if (
+        await ConversationClassification.findOne({
+          where: { conversationId: conversation.id },
+        })
+      ) {
+        console.log("already classified", conversation.id);
+        continue;
+      }
+
+      const renderedConversationString = renderedConversation
+        .map((message) => `${message.username}: ${message.content}`)
+        .join("\n");
+      const result = await classifyConversation(renderedConversationString);
+      console.log(
+        `[%s] [%s]\n\n--------------\n\n`,
+        renderedConversationString.substring(0, 250),
+        result
+      );
+      if (result && isMessageClassification(result)) {
+        await ConversationClassification.upsert({
+          messageClass: result,
+          conversationId: conversation.id,
+        });
+      } else {
+        console.log("could not classify message", conversation.id);
       }
     }
   }

front/lib/models/user_message_classification.ts → front/lib/models/conversation_classification.ts

@@ -8,22 +8,22 @@ import type {
 import { DataTypes, Model } from "sequelize";
 
 import { front_sequelize } from "@app/lib/databases";
-import { UserMessage } from "@app/lib/models/assistant/conversation";
+import { Conversation } from "@app/lib/models/assistant/conversation";
 
-export class UserMessageClassification extends Model<
-  InferAttributes<UserMessageClassification>,
-  InferCreationAttributes<UserMessageClassification>
+export class ConversationClassification extends Model<
+  InferAttributes<ConversationClassification>,
+  InferCreationAttributes<ConversationClassification>
 > {
   declare id: CreationOptional<number>;
   declare createdAt: CreationOptional<Date>;
   declare updatedAt: CreationOptional<Date>;
 
   declare messageClass: MESSAGE_CLASS;
 
-  declare userMessageId: ForeignKey<UserMessage["id"]> | null;
+  declare conversationId: ForeignKey<Conversation["id"]> | null;
 }
 
-UserMessageClassification.init(
+ConversationClassification.init(
   {
     id: {
       type: DataTypes.INTEGER,
@@ -49,10 +49,10 @@ UserMessageClassification.init(
   },
   {
     sequelize: front_sequelize,
-    modelName: "UserMessageClassification",
-    tableName: "user_message_classifications",
+    modelName: "ConversationClassification",
+    tableName: "conversation_classifications",
   }
 );
 
-UserMessage.hasMany(UserMessageClassification);
-UserMessageClassification.belongsTo(UserMessage);
+Conversation.hasMany(ConversationClassification);
+ConversationClassification.belongsTo(Conversation);

types/src/shared/message_classification.ts

@@ -4,7 +4,7 @@ export const MESSAGE_CLASSES = [
   "sales",
   "engineering",
   "coding",
-  "data",
+  "data-analysis",
   "customer support",
   "product",
   "product marketing",
@@ -14,6 +14,12 @@ export const MESSAGE_CLASSES = [
   "hiring",
   "ops",
   "security compliance",
+  "summarization",
+  "translation",
+  "joking",
+  "how-dust-works",
+  "writing-compagnion",
+  "search",
   "unknown",
 ] as const;