@ainetwork/adk-provider-memory-mongodb 0.4.1 → 0.5.1

package/dist/index.cjs

@@ -35,7 +35,7 @@ __export(index_exports, {
 module.exports = __toCommonJS(index_exports);
 
 // implements/base.memory.ts
-var import_mongoose10 = __toESM(require("mongoose"), 1);
+var import_mongoose12 = __toESM(require("mongoose"), 1);
 var import_logger2 = require("@ainetwork/adk/utils/logger");
 
 // models/agent.model.ts
@@ -179,6 +179,11 @@ var IntentObjectSchema = new import_mongoose3.Schema(
     tags: {
       type: [String],
       required: false
+    },
+    toolChoice: {
+      type: String,
+      enum: ["auto", "required"],
+      required: false
     }
   }
 );
@@ -263,6 +268,11 @@ var ThreadObjectSchema = new import_mongoose4.Schema(
       type: Boolean,
       required: false,
       default: false
+    },
+    workflowId: {
+      type: String,
+      required: false,
+      index: true
     }
   },
   {
@@ -321,6 +331,7 @@ var MessageObjectSchema = new import_mongoose6.Schema(
     timestamps: true
   }
 );
+MessageObjectSchema.index({ threadId: 1, messageId: 1 }, { unique: true });
 var MessageModel = import_mongoose7.default.model("Message", MessageObjectSchema);
 
 // implements/thread.memory.ts
@@ -345,6 +356,7 @@ var MongoDBThread = class {
         type: thread.type,
         title: thread.title || "New thread",
         isPinned: thread.isPinned ?? false,
+        workflowId: thread.workflowId,
         messages: []
       };
       messages.forEach((message) => {
@@ -359,31 +371,39 @@ var MongoDBThread = class {
       return threadObject;
     }, `getThread(${userId}, ${threadId})`);
   }
-  async createThread(type, userId, threadId, title) {
+  async createThread(type, userId, threadId, title, workflowId) {
     return this.executeWithRetry(async () => {
       await ThreadModel.create({
         type,
         userId,
         threadId,
-        title
+        title,
+        workflowId
       });
-      return { type, userId, threadId, title, messages: [] };
+      return { type, userId, threadId, title, workflowId, messages: [] };
     }, `createThread(${userId}, ${threadId})`);
   }
   async addMessagesToThread(userId, threadId, messages) {
     return this.executeWithRetry(async () => {
-      await ThreadModel.updateOne({ threadId, userId }, { $set: {} });
-      for (const message of messages) {
-        await MessageModel.create({
-          threadId,
-          messageId: message.messageId,
-          userId,
-          role: message.role,
-          content: message.content,
-          timestamp: message.timestamp,
-          metadata: message.metadata
-        });
+      if (messages.length > 0) {
+        const messageIds = messages.map((m) => m.messageId);
+        await MessageModel.deleteMany({ threadId, userId, messageId: { $in: messageIds } });
+        await MessageModel.insertMany(
+          messages.map((message) => ({
+            threadId,
+            messageId: message.messageId,
+            userId,
+            role: message.role,
+            content: message.content,
+            timestamp: message.timestamp,
+            metadata: message.metadata
+          }))
+        );
       }
+      await ThreadModel.updateOne(
+        { threadId, userId },
+        { $set: { updatedAt: /* @__PURE__ */ new Date() } }
+      );
     }, `addMessagesToThread(${userId}, ${threadId})`);
   }
   async deleteThread(userId, threadId) {
@@ -393,10 +413,13 @@ var MongoDBThread = class {
       await ThreadModel.deleteOne({ userId, threadId }).maxTimeMS(timeout);
     }, `deleteThread(${userId}, ${threadId})`);
   }
-  async listThreads(userId) {
+  async listThreads(userId, filter) {
     return this.executeWithRetry(async () => {
       const timeout = this.getOperationTimeout();
-      const threads = await ThreadModel.find({ userId }).sort({ updatedAt: -1 }).maxTimeMS(timeout);
+      const query = { userId };
+      if (filter?.workflowId) query.workflowId = filter.workflowId;
+      if (filter?.type) query.type = filter.type;
+      const threads = await ThreadModel.find(query).sort({ updatedAt: -1 }).maxTimeMS(timeout);
       const data = threads.map((thread) => {
         return {
           type: thread.type,
@@ -404,7 +427,9 @@ var MongoDBThread = class {
           threadId: thread.threadId,
           title: thread.title,
           isPinned: thread.isPinned ?? false,
-          updatedAt: thread.updatedAt
+          workflowId: thread.workflowId,
+          createdAt: thread.createdAt?.toISOString(),
+          updatedAt: thread.updatedAt?.toISOString()
         };
       });
       return data;
@@ -421,10 +446,10 @@ var MongoDBThread = class {
   }
 };
 
-// models/workflow.model.ts
+// models/user-workflow.model.ts
 var import_mongoose8 = require("mongoose");
 var import_mongoose9 = __toESM(require("mongoose"), 1);
-var WorkflowObjectSchema = new import_mongoose8.Schema(
+var UserWorkflowObjectSchema = new import_mongoose8.Schema(
   {
     workflowId: {
       type: String,
@@ -432,84 +457,203 @@ var WorkflowObjectSchema = new import_mongoose8.Schema(
       unique: true
     },
     userId: {
-      type: String
+      type: String,
+      required: true,
+      index: true
     },
     title: {
       type: String,
       required: true
     },
     description: {
-      type: String,
-      required: true
+      type: String
     },
     active: {
       type: Boolean,
       required: true,
       default: false
     },
+    templateId: {
+      type: String
+    },
     content: {
       type: String,
       required: true
     },
     variables: {
       type: import_mongoose8.Schema.Types.Mixed
+    },
+    variableValues: {
+      type: import_mongoose8.Schema.Types.Mixed
+    },
+    schedule: {
+      type: String
+    },
+    timezone: {
+      type: String
+    },
+    lastRunAt: {
+      type: Number
+    },
+    nextRunAt: {
+      type: Number
+    },
+    lastThreadId: {
+      type: String
     }
   },
   {
     timestamps: true
   }
 );
-var WorkflowModel = import_mongoose9.default.model("Workflow", WorkflowObjectSchema);
+var UserWorkflowModel = import_mongoose9.default.model(
+  "UserWorkflow",
+  UserWorkflowObjectSchema
+);
 
-// implements/workflow.memory.ts
-var MongoDBWorkflow = class {
+// implements/user-workflow.memory.ts
+var MongoDBUserWorkflow = class {
   executeWithRetry;
   getOperationTimeout;
   constructor(executeWithRetry, getOperationTimeout) {
     this.executeWithRetry = executeWithRetry;
     this.getOperationTimeout = getOperationTimeout;
   }
-  async createWorkflow(workflow) {
+  async createUserWorkflow(workflow) {
     return this.executeWithRetry(async () => {
-      const timeout = this.getOperationTimeout();
-      const created = await WorkflowModel.create(workflow);
+      const created = await UserWorkflowModel.create(workflow);
       return created.toObject();
-    }, "createWorkflow()");
+    }, "createUserWorkflow()");
   }
-  async getWorkflow(workflowId) {
+  async getUserWorkflow(workflowId) {
     return this.executeWithRetry(async () => {
       const timeout = this.getOperationTimeout();
-      const workflow = await WorkflowModel.findOne({
-        workflowId
-      }).maxTimeMS(timeout).lean();
+      const workflow = await UserWorkflowModel.findOne({ workflowId }).maxTimeMS(timeout).lean();
       return workflow || void 0;
-    }, "getWorkflow()");
+    }, "getUserWorkflow()");
   }
-  async updateWorkflow(workflowId, updates) {
+  async updateUserWorkflow(workflowId, updates) {
     if (!updates.userId) {
-      throw new Error("userId is required for updateWorkflow");
+      throw new Error("userId is required for updateUserWorkflow");
     }
+    const { userId, workflowId: _workflowId, ...mutableUpdates } = updates;
     return this.executeWithRetry(async () => {
       const timeout = this.getOperationTimeout();
-      await WorkflowModel.updateOne(
-        { workflowId, userId: updates.userId },
-        { $set: updates }
+      await UserWorkflowModel.updateOne(
+        { workflowId, userId },
+        { $set: mutableUpdates }
       ).maxTimeMS(timeout);
-    }, "updateWorkflow()");
+    }, "updateUserWorkflow()");
   }
-  async deleteWorkflow(workflowId, userId) {
+  async deleteUserWorkflow(workflowId, userId) {
     return this.executeWithRetry(async () => {
       const timeout = this.getOperationTimeout();
-      await WorkflowModel.deleteOne({ workflowId, userId }).maxTimeMS(timeout);
-    }, "deleteWorkflow()");
+      await UserWorkflowModel.deleteOne({ workflowId, userId }).maxTimeMS(timeout);
+    }, "deleteUserWorkflow()");
   }
-  async listWorkflows(userId) {
+  async listUserWorkflows(userId) {
     return this.executeWithRetry(async () => {
       const timeout = this.getOperationTimeout();
-      const query = userId ? { $or: [{ userId }, { userId: { $exists: false } }, { userId: null }] } : { $or: [{ userId: { $exists: false } }, { userId: null }] };
-      const workflows = await WorkflowModel.find(query).maxTimeMS(timeout).lean();
+      const query = userId ? { userId } : {};
+      const workflows = await UserWorkflowModel.find(query).maxTimeMS(timeout).lean();
       return workflows;
-    }, "listWorkflows()");
+    }, "listUserWorkflows()");
+  }
+  async listActiveScheduledWorkflows() {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const workflows = await UserWorkflowModel.find({
+        active: true,
+        schedule: { $exists: true, $ne: null }
+      }).maxTimeMS(timeout).lean();
+      return workflows;
+    }, "listActiveScheduledWorkflows()");
+  }
+};
+
+// models/workflow-template.model.ts
+var import_mongoose10 = require("mongoose");
+var import_mongoose11 = __toESM(require("mongoose"), 1);
+var WorkflowTemplateObjectSchema = new import_mongoose10.Schema(
+  {
+    templateId: {
+      type: String,
+      required: true,
+      unique: true
+    },
+    title: {
+      type: String,
+      required: true
+    },
+    description: {
+      type: String,
+      required: true
+    },
+    active: {
+      type: Boolean,
+      required: true,
+      default: false
+    },
+    content: {
+      type: String,
+      required: true
+    },
+    variables: {
+      type: import_mongoose10.Schema.Types.Mixed
+    }
+  },
+  {
+    timestamps: true
+  }
+);
+var WorkflowTemplateModel = import_mongoose11.default.model(
+  "WorkflowTemplate",
+  WorkflowTemplateObjectSchema
+);
+
+// implements/workflow-template.memory.ts
+var MongoDBWorkflowTemplate = class {
+  executeWithRetry;
+  getOperationTimeout;
+  constructor(executeWithRetry, getOperationTimeout) {
+    this.executeWithRetry = executeWithRetry;
+    this.getOperationTimeout = getOperationTimeout;
+  }
+  async createTemplate(template) {
+    return this.executeWithRetry(async () => {
+      const created = await WorkflowTemplateModel.create(template);
+      return created.toObject();
+    }, "createTemplate()");
+  }
+  async getTemplate(templateId) {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const template = await WorkflowTemplateModel.findOne({ templateId }).maxTimeMS(timeout).lean();
+      return template || void 0;
+    }, "getTemplate()");
+  }
+  async updateTemplate(templateId, updates) {
+    const { templateId: _templateId, ...mutableUpdates } = updates;
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      await WorkflowTemplateModel.updateOne(
+        { templateId },
+        { $set: mutableUpdates }
+      ).maxTimeMS(timeout);
+    }, "updateTemplate()");
+  }
+  async deleteTemplate(templateId) {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      await WorkflowTemplateModel.deleteOne({ templateId }).maxTimeMS(timeout);
+    }, "deleteTemplate()");
+  }
+  async listTemplates() {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const templates = await WorkflowTemplateModel.find().maxTimeMS(timeout).lean();
+      return templates;
+    }, "listTemplates()");
   }
 };
 
@@ -530,7 +674,8 @@ var MongoDBMemory = class _MongoDBMemory {
   agentMemory;
   intentMemory;
   threadMemory;
-  workflowMemory;
+  workflowTemplateMemory;
+  userWorkflowMemory;
   constructor(config) {
     const cfg = typeof config === "string" ? { uri: config } : config;
     this.uri = cfg.uri;
@@ -541,7 +686,9 @@ var MongoDBMemory = class _MongoDBMemory {
       this.threadTTLSeconds = cfg.threadTTLSeconds;
     }
     this.connectionConfig = {
-      maxPoolSize: cfg.maxPoolSize ?? 1,
+      maxPoolSize: cfg.maxPoolSize ?? 10,
+      minPoolSize: 0,
+      maxIdleTimeMS: 3e4,
       serverSelectionTimeoutMS: cfg.serverSelectionTimeoutMS ?? 3e4,
       socketTimeoutMS: cfg.socketTimeoutMS ?? 45e3,
       connectTimeoutMS: cfg.connectTimeoutMS ?? 3e4,
@@ -566,7 +713,11 @@ var MongoDBMemory = class _MongoDBMemory {
       this.executeWithRetry.bind(this),
       this.getOperationTimeout.bind(this)
     );
-    this.workflowMemory = new MongoDBWorkflow(
+    this.workflowTemplateMemory = new MongoDBWorkflowTemplate(
+      this.executeWithRetry.bind(this),
+      this.getOperationTimeout.bind(this)
+    );
+    this.userWorkflowMemory = new MongoDBUserWorkflow(
       this.executeWithRetry.bind(this),
       this.getOperationTimeout.bind(this)
     );
@@ -580,29 +731,32 @@ var MongoDBMemory = class _MongoDBMemory {
   getIntentMemory() {
     return this.intentMemory;
   }
-  getWorkflowMemory() {
-    return this.workflowMemory;
+  getWorkflowTemplateMemory() {
+    return this.workflowTemplateMemory;
+  }
+  getUserWorkflowMemory() {
+    return this.userWorkflowMemory;
   }
   setupMongooseEventListeners() {
     if (this.eventListenersSetup) return;
     this.eventListenersSetup = true;
-    import_mongoose10.default.connection.on("connected", () => {
+    import_mongoose12.default.connection.on("connected", () => {
       this.connected = true;
       this.reconnectAttempts = 0;
       this.reconnecting = false;
       import_logger2.loggers.agent.info("MongoDB connected successfully");
     });
-    import_mongoose10.default.connection.on("disconnected", () => {
+    import_mongoose12.default.connection.on("disconnected", () => {
       this.connected = false;
       import_logger2.loggers.agent.warn("MongoDB disconnected");
       this.handleDisconnection();
     });
-    import_mongoose10.default.connection.on("error", (error) => {
+    import_mongoose12.default.connection.on("error", (error) => {
       this.connected = false;
       import_logger2.loggers.agent.error("MongoDB connection error:", error);
       this.handleDisconnection();
     });
-    import_mongoose10.default.connection.on("reconnected", () => {
+    import_mongoose12.default.connection.on("reconnected", () => {
       this.connected = true;
       this.reconnectAttempts = 0;
       this.reconnecting = false;
@@ -614,13 +768,13 @@ var MongoDBMemory = class _MongoDBMemory {
       return;
     }
     this.reconnecting = true;
-    while (this.reconnectAttempts < this.maxReconnectAttempts && !this.isConnected) {
+    while (this.reconnectAttempts < this.maxReconnectAttempts && !this.connected) {
       this.reconnectAttempts++;
       import_logger2.loggers.agent.info(
         `Attempting to reconnect to MongoDB (${this.reconnectAttempts}/${this.maxReconnectAttempts})...`
       );
       try {
-        await import_mongoose10.default.connect(this.uri, this.connectionConfig);
+        await import_mongoose12.default.connect(this.uri, this.connectionConfig);
         this.connected = true;
         this.reconnectAttempts = 0;
         this.reconnecting = false;
@@ -639,7 +793,7 @@ var MongoDBMemory = class _MongoDBMemory {
       }
     }
     this.reconnecting = false;
-    if (!this.isConnected) {
+    if (!this.connected) {
       import_logger2.loggers.agent.error(
         `Failed to reconnect to MongoDB after ${this.maxReconnectAttempts} attempts`
       );
@@ -650,7 +804,7 @@ var MongoDBMemory = class _MongoDBMemory {
       return;
     }
     try {
-      await import_mongoose10.default.connect(this.uri, this.connectionConfig);
+      await import_mongoose12.default.connect(this.uri, this.connectionConfig);
       this.connected = true;
       this.reconnectAttempts = 0;
       await this.setupTTLIndex();
@@ -661,7 +815,7 @@ var MongoDBMemory = class _MongoDBMemory {
     }
   }
   async disconnect() {
-    if (!this.isConnected) {
+    if (!this.connected) {
       return;
     }
     try {
@@ -669,7 +823,7 @@ var MongoDBMemory = class _MongoDBMemory {
         clearInterval(this.orphanCleanupTimer);
         this.orphanCleanupTimer = void 0;
       }
-      await import_mongoose10.default.disconnect();
+      await import_mongoose12.default.disconnect();
       this.connected = false;
     } catch (error) {
       import_logger2.loggers.agent.error("Failed to disconnect from MongoDB:", error);
@@ -680,7 +834,7 @@ var MongoDBMemory = class _MongoDBMemory {
     return this.connected;
   }
   async ensureConnection() {
-    if (!this.isConnected && !this.reconnecting) {
+    if (!this.connected && !this.reconnecting) {
       await this.connect();
     }
     const maxWaitTime = 3e4;
@@ -688,14 +842,14 @@ var MongoDBMemory = class _MongoDBMemory {
     while (this.reconnecting && Date.now() - startTime < maxWaitTime) {
       await new Promise((resolve) => setTimeout(resolve, 100));
     }
-    if (!this.isConnected) {
+    if (!this.connected) {
       throw new Error("MongoDB is not connected and reconnection failed");
     }
   }
   async setupTTLIndex() {
     if (this.threadTTLSeconds === void 0) return;
     try {
-      const db = import_mongoose10.default.connection.db;
+      const db = import_mongoose12.default.connection.db;
       if (!db) return;
       const collection = db.collection("threads");
       const indexes = await collection.indexes();
@@ -735,7 +889,7 @@ var MongoDBMemory = class _MongoDBMemory {
   async cleanupOrphanedMessages() {
     if (!this.connected) return;
     try {
-      const db = import_mongoose10.default.connection.db;
+      const db = import_mongoose12.default.connection.db;
       if (!db) return;
       const existingThreadIds = await db.collection("threads").distinct("threadId");
       const result = await MessageModel.deleteMany({
@@ -767,6 +921,25 @@ var MongoDBMemory = class _MongoDBMemory {
         import_logger2.loggers.agent.error(`${operationName} exceeded time limit`);
         throw error;
       }
+      if (error.code === 261 || error.codeName === "TooManyLogicalSessions") {
+        import_logger2.loggers.agent.warn(
+          `${operationName} failed due to too many sessions, disconnecting to release sessions...`
+        );
+        try {
+          await import_mongoose12.default.disconnect();
+          this.connected = false;
+        } catch (disconnectError) {
+          import_logger2.loggers.agent.error("Failed to disconnect during session cleanup:", disconnectError);
+        }
+        await new Promise((resolve) => setTimeout(resolve, 5e3));
+        await this.ensureConnection();
+        try {
+          return await operation();
+        } catch (retryError) {
+          import_logger2.loggers.agent.error(`${operationName} failed after session cleanup retry:`, retryError);
+          throw retryError;
+        }
+      }
       if (error.name === "MongoNetworkError" || error.name === "MongoServerError" || error.message?.includes("connection") || error.message?.includes("disconnect")) {
         import_logger2.loggers.agent.warn(
           `${operationName} failed due to connection issue, attempting reconnection...`