@ainetwork/adk-provider-memory-mongodb 0.2.3 → 0.3.0

package/dist/chunk-OWGGE5EW.js

@@ -0,0 +1,17 @@
+// models/agent.model.ts
+import { Schema } from "mongoose";
+import mongoose from "mongoose";
+var AgentObjectSchema = new Schema(
+  {
+    agent_prompt: {
+      type: String
+    }
+  }
+);
+var AgentModel = mongoose.model("Agent", AgentObjectSchema);
+
+export {
+  AgentObjectSchema,
+  AgentModel
+};
+//# sourceMappingURL=chunk-OWGGE5EW.js.map

package/dist/chunk-OWGGE5EW.js.map

@@ -0,0 +1 @@
+{"version":3,"sources":["../models/agent.model.ts"],"sourcesContent":["import { type Document, Schema } from \"mongoose\";\nimport mongoose from \"mongoose\";\n\nexport const AgentObjectSchema = new Schema(\n\t{\n\t\tagent_prompt: {\n\t\t\ttype: String,\n\t\t},\n\t},\n);\n\nexport interface AgentDocument extends Document {\n\tagent_prompt: string;\n}\n\nexport const AgentModel = mongoose.model<AgentDocument>(\"Agent\", AgentObjectSchema);\n"],"mappings":";AAAA,SAAwB,cAAc;AACtC,OAAO,cAAc;AAEd,IAAM,oBAAoB,IAAI;AAAA,EACpC;AAAA,IACC,cAAc;AAAA,MACb,MAAM;AAAA,IACP;AAAA,EACD;AACD;AAMO,IAAM,aAAa,SAAS,MAAqB,SAAS,iBAAiB;","names":[]}

package/dist/index.cjs

@@ -30,62 +30,147 @@ var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: tru
 // index.ts
 var index_exports = {};
 __export(index_exports, {
-  MongoDBIntent: () => MongoDBIntent,
-  MongoDBThread: () => MongoDBThread
+  MongoDBMemory: () => MongoDBMemory
 });
 module.exports = __toCommonJS(index_exports);
 
 // implements/base.memory.ts
-var import_mongoose = __toESM(require("mongoose"), 1);
-var import_logger = require("@ainetwork/adk/utils/logger");
-var MongoDBMemory = class {
-  _isConnected = false;
-  _uri;
-  constructor(uri) {
-    this._uri = uri;
-  }
-  async connect() {
-    if (this._isConnected) {
-      return;
-    }
-    try {
-      await import_mongoose.default.connect(this._uri, {
-        maxPoolSize: 1,
-        serverSelectionTimeoutMS: 3e4,
-        socketTimeoutMS: 45e3,
-        connectTimeoutMS: 3e4,
-        bufferCommands: false
-      });
-      this._isConnected = true;
-      import_logger.loggers.agent.info("MongoDB connected successfully");
-    } catch (error) {
-      import_logger.loggers.agent.error("Failed to connect to MongoDB:", error);
-      throw error;
+var import_mongoose8 = __toESM(require("mongoose"), 1);
+var import_logger2 = require("@ainetwork/adk/utils/logger");
+
+// models/agent.model.ts
+var import_mongoose = require("mongoose");
+var import_mongoose2 = __toESM(require("mongoose"), 1);
+var AgentObjectSchema = new import_mongoose.Schema(
+  {
+    agent_prompt: {
+      type: String
     }
   }
-  async disconnect() {
-    if (!this.isConnected) {
-      return;
-    }
-    try {
-      await import_mongoose.default.disconnect();
-      this._isConnected = false;
-      import_logger.loggers.agent.info("MongoDB disconnected successfully");
-    } catch (error) {
-      import_logger.loggers.agent.error("Failed to disconnect from MongoDB:", error);
-      throw error;
+);
+var AgentModel = import_mongoose2.default.model("Agent", AgentObjectSchema);
+
+// implements/agent.memory.ts
+var MongoDBAgent = class {
+  executeWithRetry;
+  getOperationTimeout;
+  constructor(executeWithRetry, getOperationTimeout) {
+    this.executeWithRetry = executeWithRetry;
+    this.getOperationTimeout = getOperationTimeout;
+  }
+  async getAgentPrompt() {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const metadata = await AgentModel.findOne({
+        id: "agent_metadata"
+      }).maxTimeMS(timeout).lean();
+      return metadata?.agent_prompt || "";
+    }, "getAgentPrompt()");
+  }
+  async updateAgentPrompt(prompt) {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      await AgentModel.updateOne({
+        id: "agent_metadata"
+      }, { "agent_prompt": prompt }).maxTimeMS(timeout);
+    }, "updateAgentPrompt()");
+  }
+};
+
+// models/intent.model.ts
+var import_mongoose3 = __toESM(require("mongoose"), 1);
+var IntentObjectSchema = new import_mongoose3.Schema(
+  {
+    id: {
+      type: String,
+      required: true,
+      index: true,
+      unique: true
+    },
+    name: {
+      type: String,
+      required: true,
+      index: true
+    },
+    description: {
+      type: String,
+      required: true
+    },
+    prompt: {
+      type: String,
+      required: false
+    },
+    status: {
+      type: String,
+      required: true
+    },
+    triggeringSentences: {
+      type: [String],
+      required: false
+    },
+    tags: {
+      type: [String],
+      required: false
     }
   }
-  isConnected() {
-    return this._isConnected;
+);
+var IntentModel = import_mongoose3.default.model("Intent", IntentObjectSchema);
+
+// implements/intent.memory.ts
+var MongoDBIntent = class {
+  executeWithRetry;
+  getOperationTimeout;
+  constructor(executeWithRetry, getOperationTimeout) {
+    this.executeWithRetry = executeWithRetry;
+    this.getOperationTimeout = getOperationTimeout;
+  }
+  async getIntent(intentId) {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const intent = await IntentModel.findOne({ id: intentId }).maxTimeMS(timeout).lean();
+      return intent || void 0;
+    }, `getIntent(${intentId})`);
+  }
+  async getIntentByName(intentName) {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const intent = await IntentModel.findOne({ name: intentName }).maxTimeMS(timeout).lean();
+      return intent || void 0;
+    }, `getIntentByName(${intentName})`);
+  }
+  async saveIntent(intent) {
+    return this.executeWithRetry(async () => {
+      await IntentModel.create(intent);
+    }, `saveIntent(${intent.id})`);
+  }
+  async updateIntent(intentId, intent) {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      await IntentModel.updateOne({
+        id: intentId
+      }, intent).maxTimeMS(timeout);
+    }, `updateIntent(${intentId})`);
+  }
+  async deleteIntent(intentId) {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      await IntentModel.deleteOne({ id: intentId }).maxTimeMS(timeout);
+    }, `deleteIntent(${intentId})`);
+  }
+  async listIntents() {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const intents = await IntentModel.find().maxTimeMS(timeout).lean();
+      return intents;
+    }, `listIntents()`);
   }
 };
 
 // models/threads.model.ts
 var import_memory = require("@ainetwork/adk/types/memory");
-var import_mongoose2 = require("mongoose");
-var import_mongoose3 = __toESM(require("mongoose"), 1);
-var ThreadObjectSchema = new import_mongoose2.Schema(
+var import_mongoose4 = require("mongoose");
+var import_mongoose5 = __toESM(require("mongoose"), 1);
+var ThreadObjectSchema = new import_mongoose4.Schema(
   {
     type: {
       type: String,
@@ -116,20 +201,20 @@ var ThreadObjectSchema = new import_mongoose2.Schema(
     }
   }
 );
-var ThreadModel = import_mongoose3.default.model("Thread", ThreadObjectSchema);
+var ThreadModel = import_mongoose5.default.model("Thread", ThreadObjectSchema);
 
 // models/messages.model.ts
 var import_memory2 = require("@ainetwork/adk/types/memory");
-var import_mongoose4 = require("mongoose");
-var import_mongoose5 = __toESM(require("mongoose"), 1);
-var MessageContentObjectSchema = new import_mongoose4.Schema(
+var import_mongoose6 = require("mongoose");
+var import_mongoose7 = __toESM(require("mongoose"), 1);
+var MessageContentObjectSchema = new import_mongoose6.Schema(
   {
     type: { type: String, required: true },
-    parts: { type: [import_mongoose4.Schema.Types.Mixed], required: true }
+    parts: { type: [import_mongoose6.Schema.Types.Mixed], required: true }
   },
   { _id: false }
 );
-var MessageObjectSchema = new import_mongoose4.Schema(
+var MessageObjectSchema = new import_mongoose6.Schema(
   {
     messageId: {
       type: String,
@@ -160,167 +245,306 @@ var MessageObjectSchema = new import_mongoose4.Schema(
       required: true
     },
     metadata: {
-      type: import_mongoose4.Schema.Types.Mixed,
+      type: import_mongoose6.Schema.Types.Mixed,
       default: {}
     }
   }
 );
-var MessageModel = import_mongoose5.default.model("Message", MessageObjectSchema);
+var MessageModel = import_mongoose7.default.model("Message", MessageObjectSchema);
 
 // implements/thread.memory.ts
-var import_logger2 = require("@ainetwork/adk/utils/logger");
-var MongoDBThread = class extends MongoDBMemory {
-  constructor(uri) {
-    super(uri);
+var import_logger = require("@ainetwork/adk/utils/logger");
+var MongoDBThread = class {
+  executeWithRetry;
+  getOperationTimeout;
+  constructor(executeWithRetry, getOperationTimeout) {
+    this.executeWithRetry = executeWithRetry;
+    this.getOperationTimeout = getOperationTimeout;
   }
   async getThread(userId, threadId) {
-    const thread = await ThreadModel.findOne({ threadId, userId });
-    const messages = await MessageModel.find({ threadId, userId }).sort({
-      timestamp: 1
-    });
-    if (!thread) return void 0;
-    import_logger2.loggers.agent.debug(`Found ${messages.length} messages for thread ${threadId}`);
-    const threadObject = {
-      threadId: thread.threadId,
-      userId: thread.userId,
-      type: thread.type,
-      title: thread.title || "New thread",
-      messages: []
-    };
-    messages.forEach((message) => {
-      threadObject.messages.push({
-        messageId: message.messageId,
-        role: message.role,
-        content: message.content,
-        timestamp: message.timestamp,
-        metadata: message.metadata
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const thread = await ThreadModel.findOne({ threadId, userId }).maxTimeMS(timeout);
+      const messages = await MessageModel.find({ threadId, userId }).sort({ timestamp: 1 }).maxTimeMS(timeout);
+      if (!thread) return void 0;
+      import_logger.loggers.agent.debug(`Found ${messages.length} messages for thread ${threadId}`);
+      const threadObject = {
+        threadId: thread.threadId,
+        userId: thread.userId,
+        type: thread.type,
+        title: thread.title || "New thread",
+        messages: []
+      };
+      messages.forEach((message) => {
+        threadObject.messages.push({
+          messageId: message.messageId,
+          role: message.role,
+          content: message.content,
+          timestamp: message.timestamp,
+          metadata: message.metadata
+        });
       });
-    });
-    return threadObject;
+      return threadObject;
+    }, `getThread(${userId}, ${threadId})`);
   }
   async createThread(type, userId, threadId, title) {
-    const now = Date.now();
-    await ThreadModel.create({
-      type,
-      userId,
-      threadId,
-      title,
-      updated_at: now,
-      created_at: now
-    });
-    return { type, userId, threadId, title, messages: [] };
+    return this.executeWithRetry(async () => {
+      const now = Date.now();
+      await ThreadModel.create({
+        type,
+        userId,
+        threadId,
+        title,
+        updated_at: now,
+        created_at: now
+      });
+      return { type, userId, threadId, title, messages: [] };
+    }, `createThread(${userId}, ${threadId})`);
   }
   async addMessagesToThread(userId, threadId, messages) {
-    await ThreadModel.updateOne({ threadId, userId }, {
-      updated_at: Date.now()
-    });
-    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
+    return this.executeWithRetry(async () => {
+      await ThreadModel.updateOne({ threadId, userId }, {
+        updated_at: Date.now()
       });
-    }
+      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
+        });
+      }
+    }, `addMessagesToThread(${userId}, ${threadId})`);
   }
   async deleteThread(userId, threadId) {
-    const messages = await MessageModel.find({ userId, threadId }).sort({
-      timestamp: 1
-    });
-    messages?.forEach((message) => {
-      message.deleteOne();
-    });
-    const thread = await ThreadModel.findOne({ userId, threadId });
-    thread?.deleteOne();
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const messages = await MessageModel.find({ userId, threadId }).sort({ timestamp: 1 }).maxTimeMS(timeout);
+      messages?.forEach((message) => {
+        message.deleteOne();
+      });
+      const thread = await ThreadModel.findOne({ userId, threadId }).maxTimeMS(timeout);
+      thread?.deleteOne();
+    }, `deleteThread(${userId}, ${threadId})`);
   }
   async listThreads(userId) {
-    const threads = await ThreadModel.find({ userId }).sort({
-      updated_at: -1
-    });
-    const data = threads.map((thread) => {
-      return {
-        type: thread.type,
-        userId,
-        threadId: thread.threadId,
-        title: thread.title,
-        updatedAt: thread.updated_at
-      };
-    });
-    return data;
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const threads = await ThreadModel.find({ userId }).sort({ updated_at: -1 }).maxTimeMS(timeout);
+      const data = threads.map((thread) => {
+        return {
+          type: thread.type,
+          userId,
+          threadId: thread.threadId,
+          title: thread.title,
+          updatedAt: thread.updated_at
+        };
+      });
+      return data;
+    }, `listThreads(${userId})`);
   }
 };
 
-// models/intent.model.ts
-var import_mongoose6 = __toESM(require("mongoose"), 1);
-var IntentObjectSchema = new import_mongoose6.Schema(
-  {
-    id: {
-      type: String,
-      required: true,
-      index: true,
-      unique: true
-    },
-    name: {
-      type: String,
-      required: true,
-      index: true
-    },
-    description: {
-      type: String,
-      required: true
-    },
-    prompt: {
-      type: String,
-      required: false
-    },
-    status: {
-      type: String,
-      required: true
-    },
-    triggeringSentences: {
-      type: [String],
-      required: false
-    },
-    tags: {
-      type: [String],
-      required: false
+// implements/base.memory.ts
+var MongoDBMemory = class _MongoDBMemory {
+  static instance;
+  uri;
+  connected = false;
+  reconnectAttempts = 0;
+  maxReconnectAttempts;
+  reconnectInterval;
+  reconnecting = false;
+  connectionConfig;
+  eventListenersSetup = false;
+  operationTimeoutMS;
+  agentMemory;
+  intentMemory;
+  threadMemory;
+  constructor(config) {
+    const cfg = typeof config === "string" ? { uri: config } : config;
+    this.uri = cfg.uri;
+    this.maxReconnectAttempts = cfg.maxReconnectAttempts ?? 5;
+    this.reconnectInterval = cfg.reconnectInterval ?? 5e3;
+    this.operationTimeoutMS = cfg.operationTimeoutMS ?? 1e4;
+    this.connectionConfig = {
+      maxPoolSize: cfg.maxPoolSize ?? 1,
+      serverSelectionTimeoutMS: cfg.serverSelectionTimeoutMS ?? 3e4,
+      socketTimeoutMS: cfg.socketTimeoutMS ?? 45e3,
+      connectTimeoutMS: cfg.connectTimeoutMS ?? 3e4,
+      bufferCommands: false
+    };
+    if (!_MongoDBMemory.instance) {
+      _MongoDBMemory.instance = this;
+      this.setupMongooseEventListeners();
+    } else {
+      this.connected = _MongoDBMemory.instance.connected;
+      this.operationTimeoutMS = _MongoDBMemory.instance.operationTimeoutMS;
     }
+    this.agentMemory = new MongoDBAgent(
+      this.executeWithRetry.bind(this),
+      this.getOperationTimeout.bind(this)
+    );
+    this.threadMemory = new MongoDBThread(
+      this.executeWithRetry.bind(this),
+      this.getOperationTimeout.bind(this)
+    );
+    this.intentMemory = new MongoDBIntent(
+      this.executeWithRetry.bind(this),
+      this.getOperationTimeout.bind(this)
+    );
   }
-);
-var IntentModel = import_mongoose6.default.model("Intent", IntentObjectSchema);
-
-// implements/intent.memory.ts
-var MongoDBIntent = class extends MongoDBMemory {
-  async getIntent(intentId) {
-    const intent = await IntentModel.findOne({ id: intentId }).lean();
-    return intent || void 0;
+  getAgentMemory() {
+    return this.agentMemory;
   }
-  async getIntentByName(intentName) {
-    const intent = await IntentModel.findOne({ name: intentName }).lean();
-    return intent || void 0;
+  getThreadMemory() {
+    return this.threadMemory;
   }
-  async saveIntent(intent) {
-    await IntentModel.create(intent);
+  getIntentMemory() {
+    return this.intentMemory;
   }
-  async updateIntent(intentId, intent) {
-    await IntentModel.updateOne({
-      id: intentId
-    }, intent);
+  setupMongooseEventListeners() {
+    if (this.eventListenersSetup) return;
+    this.eventListenersSetup = true;
+    import_mongoose8.default.connection.on("connected", () => {
+      this.connected = true;
+      this.reconnectAttempts = 0;
+      this.reconnecting = false;
+      import_logger2.loggers.agent.info("MongoDB connected successfully");
+    });
+    import_mongoose8.default.connection.on("disconnected", () => {
+      this.connected = false;
+      import_logger2.loggers.agent.warn("MongoDB disconnected");
+      this.handleDisconnection();
+    });
+    import_mongoose8.default.connection.on("error", (error) => {
+      this.connected = false;
+      import_logger2.loggers.agent.error("MongoDB connection error:", error);
+      this.handleDisconnection();
+    });
+    import_mongoose8.default.connection.on("reconnected", () => {
+      this.connected = true;
+      this.reconnectAttempts = 0;
+      this.reconnecting = false;
+      import_logger2.loggers.agent.info("MongoDB reconnected successfully");
+    });
   }
-  async deleteIntent(intentId) {
-    await IntentModel.deleteOne({ id: intentId });
+  async handleDisconnection() {
+    if (this.reconnecting) {
+      return;
+    }
+    this.reconnecting = true;
+    while (this.reconnectAttempts < this.maxReconnectAttempts && !this.isConnected) {
+      this.reconnectAttempts++;
+      import_logger2.loggers.agent.info(
+        `Attempting to reconnect to MongoDB (${this.reconnectAttempts}/${this.maxReconnectAttempts})...`
+      );
+      try {
+        await import_mongoose8.default.connect(this.uri, this.connectionConfig);
+        this.connected = true;
+        this.reconnectAttempts = 0;
+        this.reconnecting = false;
+        import_logger2.loggers.agent.info("MongoDB reconnection successful");
+        return;
+      } catch (error) {
+        import_logger2.loggers.agent.error(
+          `Reconnection attempt ${this.reconnectAttempts} failed:`,
+          error
+        );
+        if (this.reconnectAttempts < this.maxReconnectAttempts) {
+          await new Promise(
+            (resolve) => setTimeout(resolve, this.reconnectInterval)
+          );
+        }
+      }
+    }
+    this.reconnecting = false;
+    if (!this.isConnected) {
+      import_logger2.loggers.agent.error(
+        `Failed to reconnect to MongoDB after ${this.maxReconnectAttempts} attempts`
+      );
+    }
   }
-  async listIntents() {
-    const intents = await IntentModel.find().lean();
-    return intents;
+  async connect() {
+    if (this.connected) {
+      return;
+    }
+    try {
+      await import_mongoose8.default.connect(this.uri, this.connectionConfig);
+      this.connected = true;
+      this.reconnectAttempts = 0;
+    } catch (error) {
+      import_logger2.loggers.agent.error("Failed to connect to MongoDB:", error);
+      throw error;
+    }
+  }
+  async disconnect() {
+    if (!this.isConnected) {
+      return;
+    }
+    try {
+      await import_mongoose8.default.disconnect();
+      this.connected = false;
+    } catch (error) {
+      import_logger2.loggers.agent.error("Failed to disconnect from MongoDB:", error);
+      throw error;
+    }
+  }
+  isConnected() {
+    return this.connected;
+  }
+  async ensureConnection() {
+    if (!this.isConnected && !this.reconnecting) {
+      await this.connect();
+    }
+    const maxWaitTime = 3e4;
+    const startTime = Date.now();
+    while (this.reconnecting && Date.now() - startTime < maxWaitTime) {
+      await new Promise((resolve) => setTimeout(resolve, 100));
+    }
+    if (!this.isConnected) {
+      throw new Error("MongoDB is not connected and reconnection failed");
+    }
+  }
+  /**
+   * Get the operation timeout in milliseconds
+   */
+  getOperationTimeout() {
+    return this.operationTimeoutMS;
+  }
+  /**
+   * Execute a database operation with automatic retry on connection errors
+   * Note: Use mongoose's maxTimeMS option in queries for timeout control
+   */
+  async executeWithRetry(operation, operationName = "Database operation") {
+    await this.ensureConnection();
+    try {
+      return await operation();
+    } catch (error) {
+      if (error.code === 50 || error.message?.includes("operation exceeded time limit")) {
+        import_logger2.loggers.agent.error(`${operationName} exceeded time limit`);
+        throw error;
+      }
+      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...`
+        );
+        await this.ensureConnection();
+        try {
+          return await operation();
+        } catch (retryError) {
+          import_logger2.loggers.agent.error(`${operationName} failed after retry:`, retryError);
+          throw retryError;
+        }
+      }
+      throw error;
+    }
   }
 };
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {
-  MongoDBIntent,
-  MongoDBThread
+  MongoDBMemory
 });
 //# sourceMappingURL=index.cjs.map