Auto-rename thread in agent mode (#5550)

auto-rename thread in agent mode

Made-with: Cursor

frontend/src/utils/chat/agent.js

@@ -2,6 +2,7 @@ import { v4 } from "uuid";
 import { safeJsonParse } from "../request";
 import { API_BASE } from "../constants";
 import { useEffect, useState } from "react";
+import { THREAD_RENAME_EVENT } from "@/components/Sidebar/ActiveWorkspaces/ThreadContainer";
 
 export const AGENT_SESSION_START = "agentSessionStart";
 export const AGENT_SESSION_END = "agentSessionEnd";
@@ -26,6 +27,19 @@ export default function handleSocketResponse(socket, event, setChatHistory) {
   const data = safeJsonParse(event.data, null);
   if (data === null) return;
 
+  // Handle thread rename
+  if (data.type === "rename_thread") {
+    const { slug, name } = data.content || {};
+    if (slug && name) {
+      window.dispatchEvent(
+        new CustomEvent(THREAD_RENAME_EVENT, {
+          detail: { threadSlug: slug, newName: name },
+        })
+      );
+    }
+    return;
+  }
+
   // No message type is defined then this is a generic message
   // that we need to print to the user as a system response
   if (!data.hasOwnProperty("type") && !socket.supportsAgentStreaming) {

server/endpoints/chat.js

@@ -15,7 +15,6 @@ const {
 const { writeResponseChunk } = require("../utils/helpers/chat/responses");
 const { WorkspaceThread } = require("../models/workspaceThread");
 const { User } = require("../models/user");
-const truncate = require("truncate");
 const { getModelTag } = require("./utils");
 
 function chatEndpoints(app) {
@@ -162,7 +161,7 @@ function chatEndpoints(app) {
           thread,
           workspace,
           user,
-          newName: truncate(message, 22),
+          prompt: message,
           onRename: (thread) => {
             writeResponseChunk(response, {
               action: "rename_thread",

server/models/workspaceThread.js

@@ -1,6 +1,7 @@
 const prisma = require("../utils/prisma");
 const slugifyModule = require("slugify");
 const { v4: uuidv4 } = require("uuid");
+const truncate = require("truncate");
 
 const WorkspaceThread = {
   defaultName: "Thread",
@@ -120,15 +121,15 @@ const WorkspaceThread = {
     }
   },
 
-  // Will fire on first message (included or not) for a thread and rename the thread with the newName prop.
+  // Will fire on first message (included or not) for a thread and rename the thread based on the prompt.
   autoRenameThread: async function ({
     workspace = null,
     thread = null,
     user = null,
-    newName = null,
+    prompt = null,
     onRename = null,
   }) {
-    if (!workspace || !thread || !newName) return false;
+    if (!workspace || !thread || !prompt) return false;
     if (thread.name !== this.defaultName) return false; // don't rename if already named.
 
     const { WorkspaceChats } = require("./workspaceChats");
@@ -139,7 +140,7 @@ const WorkspaceThread = {
     });
     if (chatCount !== 1) return { renamed: false, thread };
     const { thread: updatedThread } = await this.update(thread, {
-      name: newName,
+      name: truncate(prompt, 22),
     });
 
     onRename?.(updatedThread);

server/utils/agents/aibitat/plugins/chat-history.js

@@ -1,4 +1,5 @@
 const { WorkspaceChats } = require("../../../../models/workspaceChats");
+const { WorkspaceThread } = require("../../../../models/workspaceThread");
 
 /**
  * Plugin to save chat history to AnythingLLM DB.
@@ -116,6 +117,13 @@ const chatHistory = {
           threadId: invocation?.thread_id || null,
           include: true,
         });
+
+        if (!aibitat._threadRenamed) {
+          aibitat._threadRenamed = await this._autoRenameThread(
+            aibitat,
+            prompt
+          );
+        }
         this._cleanup(aibitat);
       },
       _storeSpecial: async function (
@@ -146,10 +154,43 @@ const chatHistory = {
           threadId: invocation?.thread_id || null,
           include: true,
         });
+
+        if (!aibitat._threadRenamed) {
+          aibitat._threadRenamed = await this._autoRenameThread(
+            aibitat,
+            prompt
+          );
+        }
         options?.postSave();
         this._cleanup(aibitat);
       },
 
+      _autoRenameThread: async function (aibitat, prompt) {
+        const invocation = aibitat.handlerProps.invocation;
+        if (!invocation?.thread_id) return true;
+
+        const thread = await WorkspaceThread.get({ id: invocation.thread_id });
+        if (!thread) return true;
+
+        const { Workspace } = require("../../../../models/workspace");
+        const workspace = await Workspace.get({ id: invocation.workspace_id });
+        if (!workspace) return true;
+
+        await WorkspaceThread.autoRenameThread({
+          thread,
+          workspace,
+          user: invocation.user_id ? { id: invocation.user_id } : null,
+          prompt,
+          onRename: (updatedThread) => {
+            aibitat.socket?.send("rename_thread", {
+              slug: updatedThread.slug,
+              name: updatedThread.name,
+            });
+          },
+        });
+        return true;
+      },
+
       _cleanup: function (aibitat) {
         aibitat.clearCitations?.();
         aibitat._pendingOutputs = [];