@nocobase/plugin-ai 2.4.0-alpha.3 → 2.4.0-alpha.4

package/dist/server/ai-employees/ai-employee.js

@@ -60,6 +60,7 @@ var import_attachments = require("../attachments");
 var import_frontend_tools = require("../../common/frontend-tools");
 var import_frontend_tools2 = require("../frontend-tools");
 var import_reasoning_stream_state = require("./reasoning-stream-state");
+const ABORTED_TOOL_CALL_CONTENT = "The tool call was interrupted because the conversation was aborted.";
 class AIEmployee {
   sessionId;
   from = "main-agent";
@@ -78,6 +79,7 @@ class AIEmployee {
   tools;
   inWorkflow;
   streamCached;
+  static conversationPersistenceQueues = /* @__PURE__ */ new WeakMap();
   constructor({
     ctx,
     employee,
@@ -419,6 +421,7 @@ class AIEmployee {
   async processChatStream(stream, options) {
     var _a, _b, _c, _d, _e, _f, _g, _h, _i, _j, _k, _l, _m, _n, _o, _p, _q, _r, _s, _t, _u, _v, _w, _x, _y, _z;
     const aiMessageIdMap = /* @__PURE__ */ new Map();
+    const persistedAIMessageIdMap = /* @__PURE__ */ new Map();
     const { signal, providerName, llmService, model, provider, responseMetadata, allowEmpty = false } = options;
     const reasoningState = new import_reasoning_stream_state.ReasoningStreamState();
     const stopReasoning = async (conversation) => {
@@ -432,37 +435,38 @@ class AIEmployee {
       }
     };
     let gathered;
-    signal.addEventListener("abort", async () => {
-      try {
-        await stopAllReasoning();
-        if ((gathered == null ? void 0 : gathered.type) === "ai") {
-          const values = (0, import_utils2.convertAIMessage)({
-            aiEmployee: this,
-            providerName,
-            provider,
-            llmService,
-            model,
-            aiMessage: gathered
-          });
-          if (values) {
-            values.metadata.interrupted = true;
-          }
-          await this.aiChatConversation.withTransaction(async (conversation, transaction) => {
-            const result = await conversation.addMessages(values);
-          });
-        }
-      } catch (e) {
-        this.logger.error("Fail to save message after conversation abort", gathered);
-      } finally {
-        await this.aiConversationsRepo.update({
-          values: { llmActiveState: "idle", read: true },
-          filter: {
-            sessionId: this.sessionId
+    let abortFinalization;
+    signal.addEventListener(
+      "abort",
+      () => {
+        abortFinalization = (async () => {
+          try {
+            await stopAllReasoning();
+            if ((gathered == null ? void 0 : gathered.type) === "ai") {
+              await this.finalizeAbortedAIMessage({
+                aiMessage: gathered,
+                providerName,
+                provider,
+                llmService,
+                model,
+                knownMessageId: persistedAIMessageIdMap.get(gathered.id)
+              });
+            }
+          } catch (e) {
+            this.logger.error("Fail to save message after conversation abort", gathered);
+          } finally {
+            await this.aiConversationsRepo.update({
+              values: { llmActiveState: "idle", read: true },
+              filter: {
+                sessionId: this.sessionId
+              }
+            });
+            await this.streamCached.clear();
           }
-        });
-        await this.streamCached.clear();
-      }
-    });
+        })();
+      },
+      { once: true }
+    );
     try {
       const aiEmployeeConversation = {
         sessionId: this.sessionId,
@@ -525,6 +529,7 @@ class AIEmployee {
           if (chunks.action === "AfterAIMessageSaved") {
             await this.streamCached.skipped();
             aiMessageIdMap.set(currentConversation.sessionId, chunks.body.messageId);
+            persistedAIMessageIdMap.set(chunks.body.id, chunks.body.messageId);
             const data = responseMetadata.get(chunks.body.id);
             if (data) {
               const savedMessage = await this.aiMessagesModel.findOne({
@@ -629,6 +634,7 @@ class AIEmployee {
         this.sendErrorResponse(provider.parseResponseError(err));
       }
     } finally {
+      await abortFinalization;
       if (this.from === "main-agent") {
         this.ctx.res.end();
       }
@@ -804,6 +810,170 @@ If information is missing, clearly state it in the summary.</Important>`;
     }
   }
   // === Tool calls ===
+  async withLockedConversation(callback) {
+    const persistenceQueues = AIEmployee.conversationPersistenceQueues.get(this.db) ?? /* @__PURE__ */ new Map();
+    AIEmployee.conversationPersistenceQueues.set(this.db, persistenceQueues);
+    const previous = persistenceQueues.get(this.sessionId) ?? Promise.resolve();
+    let release;
+    const current = new Promise((resolve) => {
+      release = resolve;
+    });
+    persistenceQueues.set(this.sessionId, current);
+    await previous;
+    try {
+      for (let attempt = 0; attempt < 2; attempt++) {
+        try {
+          return await this.aiChatConversation.withTransaction(async (conversation, transaction) => {
+            const lockedConversation = await this.aiConversationsModel.findOne({
+              where: { sessionId: this.sessionId },
+              transaction,
+              lock: transaction.LOCK.UPDATE
+            });
+            if (!lockedConversation) {
+              throw new Error(`AI conversation ${this.sessionId} not found`);
+            }
+            return await callback(conversation, transaction);
+          });
+        } catch (error) {
+          if (!(error instanceof import_database.UniqueConstraintError) || attempt === 1) {
+            throw error;
+          }
+        }
+      }
+      throw new Error(`Failed to persist AI conversation ${this.sessionId}`);
+    } finally {
+      release();
+      if (persistenceQueues.get(this.sessionId) === current) {
+        persistenceQueues.delete(this.sessionId);
+      }
+    }
+  }
+  async findPersistedAIMessage(transaction, langChainMessageId, knownMessageId) {
+    var _a;
+    if (knownMessageId) {
+      const knownMessage = await this.aiMessagesModel.findOne({
+        where: { sessionId: this.sessionId, messageId: knownMessageId },
+        transaction
+      });
+      if (((_a = knownMessage == null ? void 0 : knownMessage.get("metadata")) == null ? void 0 : _a.id) === langChainMessageId) {
+        return knownMessage;
+      }
+    }
+    const messages = await this.aiMessagesModel.findAll({
+      where: { sessionId: this.sessionId, role: this.employee.username },
+      order: [["messageId", "DESC"]],
+      transaction
+    });
+    return messages.find((message) => {
+      var _a2;
+      return ((_a2 = message.get("metadata")) == null ? void 0 : _a2.id) === langChainMessageId;
+    });
+  }
+  async persistAIMessageInTransaction(conversation, transaction, options) {
+    const existingMessage = await this.findPersistedAIMessage(
+      transaction,
+      options.langChainMessageId,
+      options.knownMessageId
+    );
+    if (existingMessage) {
+      const message2 = existingMessage.toJSON();
+      const initializedToolCalls2 = await this.aiToolMessagesModel.findAll({
+        where: { sessionId: this.sessionId, messageId: message2.messageId },
+        transaction
+      });
+      return { message: message2, initializedToolCalls: initializedToolCalls2, created: false };
+    }
+    const message = await conversation.addMessages(options.values);
+    const initializedToolCalls = options.toolCalls.length ? await this.initToolCall(transaction, message.messageId, options.toolCalls) : [];
+    return { message, initializedToolCalls, created: true };
+  }
+  async persistAIMessage(options) {
+    return await this.withLockedConversation(async (conversation, transaction) => {
+      return await this.persistAIMessageInTransaction(conversation, transaction, options);
+    });
+  }
+  async finalizeAbortedAIMessage({
+    aiMessage,
+    providerName,
+    provider,
+    llmService,
+    model,
+    knownMessageId
+  }) {
+    const values = (0, import_utils2.convertAIMessage)({
+      aiEmployee: this,
+      providerName,
+      provider,
+      llmService,
+      model,
+      aiMessage
+    });
+    if (!values) {
+      return;
+    }
+    values.metadata = { ...values.metadata, interrupted: true };
+    const toolCalls = aiMessage.tool_calls ?? [];
+    return await this.withLockedConversation(async (conversation, transaction) => {
+      var _a;
+      const result = await this.persistAIMessageInTransaction(conversation, transaction, {
+        values,
+        langChainMessageId: aiMessage.id,
+        toolCalls,
+        knownMessageId
+      });
+      const unfinishedToolCalls = result.initializedToolCalls.map((toolCall) => toolCall.toJSON()).filter((toolCall) => toolCall.invokeStatus !== "confirmed");
+      if (!result.created && (!result.initializedToolCalls.length || unfinishedToolCalls.length) && !((_a = result.message.metadata) == null ? void 0 : _a.interrupted)) {
+        const metadata = { ...result.message.metadata, interrupted: true };
+        await this.aiMessagesModel.update(
+          { metadata },
+          { where: { sessionId: this.sessionId, messageId: result.message.messageId }, transaction }
+        );
+        result.message.metadata = metadata;
+      }
+      if (!unfinishedToolCalls.length) {
+        return result;
+      }
+      const persistedToolCalls = result.message.toolCalls ?? toolCalls;
+      const toolCallMap = new Map(persistedToolCalls.map((toolCall) => [toolCall.id, toolCall]));
+      const now = /* @__PURE__ */ new Date();
+      await conversation.addMessages(
+        unfinishedToolCalls.map((toolCall) => ({
+          role: "tool",
+          content: { type: "text", content: ABORTED_TOOL_CALL_CONTENT },
+          metadata: {
+            id: `aborted-tool:${aiMessage.id}:${toolCall.toolCallId}`,
+            model,
+            provider: providerName,
+            llmService,
+            messageId: result.message.messageId,
+            toolCallId: toolCall.toolCallId,
+            toolName: toolCall.toolName,
+            toolCall: toolCallMap.get(toolCall.toolCallId),
+            autoCall: toolCall.auto
+          }
+        }))
+      );
+      for (const toolCall of unfinishedToolCalls) {
+        await this.aiToolMessagesModel.update(
+          {
+            invokeStatus: "confirmed",
+            status: "error",
+            content: ABORTED_TOOL_CALL_CONTENT,
+            invokeStartTime: toolCall.invokeStartTime ?? now,
+            invokeEndTime: now
+          },
+          {
+            where: {
+              id: toolCall.id,
+              invokeStatus: { [import_database.Op.ne]: "confirmed" }
+            },
+            transaction
+          }
+        );
+      }
+      return result;
+    });
+  }
   async initToolCall(transaction, messageId, toolCalls) {
     const nowTime = /* @__PURE__ */ new Date();
     const toolMap = await this.getToolsMap();
@@ -956,6 +1126,39 @@ If information is missing, clearly state it in the summary.</Important>`;
     })).map((it) => it.toJSON());
     return new Map(list.map((it) => [it.toolCallId, it]));
   }
+  async getToolCallResults(toolCallIds) {
+    if (!toolCallIds.length) {
+      return /* @__PURE__ */ new Map();
+    }
+    const list = (await this.aiToolMessagesModel.findAll({
+      where: {
+        sessionId: this.sessionId,
+        toolCallId: {
+          [import_database.Op.in]: toolCallIds
+        }
+      }
+    })).map((item) => item.toJSON());
+    const invokeStatusPriority = { confirmed: 2, done: 1 };
+    const results = /* @__PURE__ */ new Map();
+    for (const item of list) {
+      if (item.invokeStatus !== "confirmed" && item.invokeStatus !== "done") {
+        continue;
+      }
+      const existing = results.get(item.toolCallId);
+      if (!existing) {
+        results.set(item.toolCallId, item);
+        continue;
+      }
+      const priority = invokeStatusPriority[item.invokeStatus];
+      const existingPriority = invokeStatusPriority[existing.invokeStatus];
+      const updatedAt = item.updatedAt ? new Date(item.updatedAt).getTime() : 0;
+      const existingUpdatedAt = existing.updatedAt ? new Date(existing.updatedAt).getTime() : 0;
+      if (priority > existingPriority || priority === existingPriority && updatedAt > existingUpdatedAt) {
+        results.set(item.toolCallId, item);
+      }
+    }
+    return results;
+  }
   async cancelToolCall() {
     var _a;
     let messageId;
@@ -1093,7 +1296,7 @@ If information is missing, clearly state it in the summary.</Important>`;
     });
   }
   async formatMessages({ messages, provider }) {
-    var _a, _b;
+    var _a, _b, _c, _d;
     const formattedMessages = [];
     const workContextHandler = this.plugin.workContextHandler;
     const normalizedMessages = await this.normalizeMessageAttachments(messages);
@@ -1160,13 +1363,15 @@ If information is missing, clearly state it in the summary.</Important>`;
       }
       if (msg.role === "tool") {
         formattedMessages.push({
+          id: (_a = msg.metadata) == null ? void 0 : _a.id,
           role: "tool",
           content,
-          tool_call_id: (_a = msg.metadata) == null ? void 0 : _a.toolCallId
+          name: (_b = msg.metadata) == null ? void 0 : _b.toolName,
+          tool_call_id: (_c = msg.metadata) == null ? void 0 : _c.toolCallId
         });
         continue;
       }
-      const additionalKwargs = (0, import_tool_call_sanitizer.sanitizeAdditionalKwargsForToolCalls)((_b = msg.metadata) == null ? void 0 : _b.additional_kwargs, msg.toolCalls, {
+      const additionalKwargs = (0, import_tool_call_sanitizer.sanitizeAdditionalKwargsForToolCalls)((_d = msg.metadata) == null ? void 0 : _d.additional_kwargs, msg.toolCalls, {
         onDiscard: (info) => {
           var _a2;
           this.logger.warn("Discard malformed raw tool calls from AI message", {
@@ -1415,7 +1620,22 @@ If information is missing, clearly state it in the summary.</Important>`;
       (0, import_middleware.toolCallStatusMiddleware)(this),
       ...inWorkflow ? [(0, import_middleware.workflowHistoryMiddleware)(this, this.db)] : [],
       (0, import_middleware.conversationMiddleware)(this, { providerName, provider, llmService, model, messageId, agentThread }),
-      (0, import_middleware.toolCallSanitizerMiddleware)({ logger: this.logger })
+      (0, import_middleware.toolCallSanitizerMiddleware)({ logger: this.logger }),
+      (0, import_middleware.toolResultIntegrityMiddleware)({
+        sessionId: this.sessionId,
+        logger: this.logger,
+        loadToolResults: async (toolCallIds) => {
+          try {
+            return await this.getToolCallResults(toolCallIds);
+          } catch (error) {
+            this.logger.warn("Failed to load persisted tool results before model call", {
+              sessionId: this.sessionId,
+              error
+            });
+            return /* @__PURE__ */ new Map();
+          }
+        }
+      })
     ];
   }
   async getCurrentThread() {
@@ -1473,6 +1693,9 @@ If information is missing, clearly state it in the summary.</Important>`;
   get aiMessagesRepo() {
     return this.ctx.db.getRepository("aiMessages");
   }
+  get aiConversationsModel() {
+    return this.ctx.db.getModel("aiConversations");
+  }
   get aiMessagesModel() {
     return this.ctx.db.getModel("aiMessages");
   }

package/dist/server/ai-employees/ai-knowledge-base.js

@@ -51,6 +51,7 @@ module.exports = __toCommonJS(ai_knowledge_base_exports);
 var import_prompts = require("@langchain/core/prompts");
 var import_ai_feature_manager = require("../manager/ai-feature-manager");
 var import_lodash = __toESM(require("lodash"));
+var import_ai_employee_validation = require("../../common/ai-employee-validation");
 const KNOWLEDGE_BASE_RETRIEVAL_STRATEGIES = ["always", "onDemand"];
 const isRecord = (value) => value !== null && typeof value === "object" && !Array.isArray(value);
 const getRoleNames = (roles) => Array.isArray(roles) ? Array.from(
@@ -120,15 +121,21 @@ class KnowledgeBaseManager {
       return "Specified knowledge base not existed";
     }
     const { knowledgeBaseKeys = [], topK, score } = employee.knowledgeBase ?? {};
-    const promptTemplate = import_prompts.ChatPromptTemplate.fromTemplate(employee.knowledgeBasePrompt ?? "{knowledgeBaseData}");
+    const knowledgeBasePrompt = employee.knowledgeBasePrompt ?? "{knowledgeBaseData}";
     const docs = await this.plugin.features.knowledgeBase.search({ knowledgeBaseKeys, query, topK, score, roleNames });
     if (!(docs == null ? void 0 : docs.length)) {
       return "No document match in knowledge base";
     }
     const knowledgeBaseData = docs.map((doc) => buildKnowledgeBaseContent(doc.content, doc.metadata)).join("\n");
-    return import_lodash.default.isEmpty(knowledgeBaseData) ? "No document match in knowledge base" : await promptTemplate.format({
-      knowledgeBaseData
-    });
+    if (import_lodash.default.isEmpty(knowledgeBaseData)) {
+      return "No document match in knowledge base";
+    }
+    if (!(0, import_ai_employee_validation.hasKnowledgeBaseDataPlaceholder)(knowledgeBasePrompt)) {
+      return `${knowledgeBasePrompt}
+
+${knowledgeBaseData}`;
+    }
+    return import_prompts.ChatPromptTemplate.fromTemplate(knowledgeBasePrompt).format({ knowledgeBaseData });
   }
   async hasAccessibleKnowledgeBase({ employee, roleNames }) {
     var _a;

package/dist/server/ai-employees/middleware/conversation.js

@@ -179,31 +179,30 @@ const conversationMiddleware = (aiEmployee, options) => {
         const toolCalls = aiMessage.tool_calls;
         const values = convertAIMessage(aiMessage);
         if (values) {
-          await aiEmployee.aiChatConversation.withTransaction(async (conversation, transaction) => {
-            const result = await conversation.addMessages(values);
-            newState.messageId = result.messageId;
-            if (toolCalls == null ? void 0 : toolCalls.length) {
-              const toolsMap = await aiEmployee.getToolsMap();
-              const initializedToolCalls = await aiEmployee.initToolCall(
-                transaction,
-                result.messageId,
-                toolCalls
-              );
-              fillToolCall(result, toolsMap, initializedToolCalls, toolCalls);
-            }
+          const result = await aiEmployee.persistAIMessage({
+            values,
+            langChainMessageId: aiMessage.id,
+            toolCalls: toolCalls ?? []
           });
+          newState.messageId = result.message.messageId;
           if (toolCalls == null ? void 0 : toolCalls.length) {
-            (_b = runtime.writer) == null ? void 0 : _b.call(runtime, {
-              action: "initToolCalls",
-              body: { toolCalls },
+            const toolsMap = await aiEmployee.getToolsMap();
+            fillToolCall(result.message, toolsMap, result.initializedToolCalls, toolCalls);
+          }
+          if (result.created) {
+            if (toolCalls == null ? void 0 : toolCalls.length) {
+              (_b = runtime.writer) == null ? void 0 : _b.call(runtime, {
+                action: "initToolCalls",
+                body: { toolCalls },
+                currentConversation
+              });
+            }
+            (_c = runtime.writer) == null ? void 0 : _c.call(runtime, {
+              action: "AfterAIMessageSaved",
+              body: { id: aiMessage.id, messageId: newState.messageId },
               currentConversation
             });
           }
-          (_c = runtime.writer) == null ? void 0 : _c.call(runtime, {
-            action: "AfterAIMessageSaved",
-            body: { id: aiMessage.id, messageId: newState.messageId },
-            currentConversation
-          });
         }
         return newState;
       } catch (e) {

package/dist/server/ai-employees/middleware/index.d.ts

@@ -9,5 +9,6 @@
 export * from './conversation';
 export * from './skill-tools';
 export * from './tool-call-sanitizer';
+export * from './tool-result-integrity';
 export * from './tools';
 export * from './workflow-history';

package/dist/server/ai-employees/middleware/index.js

@@ -26,6 +26,7 @@ module.exports = __toCommonJS(middleware_exports);
 __reExport(middleware_exports, require("./conversation"), module.exports);
 __reExport(middleware_exports, require("./skill-tools"), module.exports);
 __reExport(middleware_exports, require("./tool-call-sanitizer"), module.exports);
+__reExport(middleware_exports, require("./tool-result-integrity"), module.exports);
 __reExport(middleware_exports, require("./tools"), module.exports);
 __reExport(middleware_exports, require("./workflow-history"), module.exports);
 // Annotate the CommonJS export names for ESM import in node:
@@ -33,6 +34,7 @@ __reExport(middleware_exports, require("./workflow-history"), module.exports);
   ...require("./conversation"),
   ...require("./skill-tools"),
   ...require("./tool-call-sanitizer"),
+  ...require("./tool-result-integrity"),
   ...require("./tools"),
   ...require("./workflow-history")
 });

package/dist/server/ai-employees/middleware/tool-result-integrity.d.ts

@@ -0,0 +1,22 @@
+/**
+ * This file is part of the NocoBase (R) project.
+ * Copyright (c) 2020-2024 NocoBase Co., Ltd.
+ * Authors: NocoBase Team.
+ *
+ * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License.
+ * For more information, please refer to: https://www.nocobase.com/agreement.
+ */
+import { BaseMessage } from '@langchain/core/messages';
+import type { AIToolMessage } from '../../types/ai-message.type';
+type ToolResultIntegrityLogger = {
+    warn: (message: string, meta?: Record<string, unknown>) => void;
+};
+export type NormalizeToolCallHistoryOptions = {
+    sessionId: string;
+    logger?: ToolResultIntegrityLogger;
+    loadToolResults: (toolCallIds: string[]) => Promise<Map<string, AIToolMessage>>;
+};
+export declare const isToolCallHistoryValid: (messages: readonly BaseMessage[]) => boolean;
+export declare const normalizeToolCallHistory: (messages: readonly BaseMessage[], options: NormalizeToolCallHistoryOptions) => Promise<BaseMessage[]>;
+export declare const toolResultIntegrityMiddleware: (options: NormalizeToolCallHistoryOptions) => import("langchain/dist/agents/middleware/types.cjs").AgentMiddleware<undefined, undefined, unknown, readonly (import("@langchain/core/dist/tools").ClientTool | import("@langchain/core/dist/tools").ServerTool)[]>;
+export {};