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

package/dist/models/agent.model.d.cts

@@ -0,0 +1,23 @@
+import mongoose, { Schema, Document } from 'mongoose';
+
+declare const AgentObjectSchema: Schema<any, mongoose.Model<any, any, any, any, any, any>, {}, {}, {}, {}, mongoose.DefaultSchemaOptions, {
+    agent_prompt?: string | null | undefined;
+}, Document<unknown, {}, mongoose.FlatRecord<{
+    agent_prompt?: string | null | undefined;
+}>, {}> & mongoose.FlatRecord<{
+    agent_prompt?: string | null | undefined;
+}> & {
+    _id: mongoose.Types.ObjectId;
+} & {
+    __v: number;
+}>;
+interface AgentDocument extends Document {
+    agent_prompt: string;
+}
+declare const AgentModel: mongoose.Model<AgentDocument, {}, {}, {}, Document<unknown, {}, AgentDocument, {}> & AgentDocument & Required<{
+    _id: unknown;
+}> & {
+    __v: number;
+}, any>;
+
+export { type AgentDocument, AgentModel, AgentObjectSchema };

package/dist/models/agent.model.d.ts

@@ -0,0 +1,23 @@
+import mongoose, { Schema, Document } from 'mongoose';
+
+declare const AgentObjectSchema: Schema<any, mongoose.Model<any, any, any, any, any, any>, {}, {}, {}, {}, mongoose.DefaultSchemaOptions, {
+    agent_prompt?: string | null | undefined;
+}, Document<unknown, {}, mongoose.FlatRecord<{
+    agent_prompt?: string | null | undefined;
+}>, {}> & mongoose.FlatRecord<{
+    agent_prompt?: string | null | undefined;
+}> & {
+    _id: mongoose.Types.ObjectId;
+} & {
+    __v: number;
+}>;
+interface AgentDocument extends Document {
+    agent_prompt: string;
+}
+declare const AgentModel: mongoose.Model<AgentDocument, {}, {}, {}, Document<unknown, {}, AgentDocument, {}> & AgentDocument & Required<{
+    _id: unknown;
+}> & {
+    __v: number;
+}, any>;
+
+export { type AgentDocument, AgentModel, AgentObjectSchema };

package/dist/models/agent.model.js

@@ -0,0 +1,9 @@
+import {
+  AgentModel,
+  AgentObjectSchema
+} from "../chunk-OWGGE5EW.js";
+export {
+  AgentModel,
+  AgentObjectSchema
+};
+//# sourceMappingURL=agent.model.js.map

package/dist/models/agent.model.js.map

@@ -0,0 +1 @@
+{"version":3,"sources":[],"sourcesContent":[],"mappings":"","names":[]}

package/implements/agent.memory.ts

@@ -0,0 +1,46 @@
+import { IAgentMemory } from "@ainetwork/adk/modules";
+import { AgentModel } from "../models/agent.model";
+
+export type ExecuteWithRetryFn = <T>(
+  operation: () => Promise<T>,
+  operationName?: string
+) => Promise<T>;
+
+export type GetOperationTimeoutFn = () => number;
+
+type AgentMetadata = {
+  agent_prompt: string;
+}
+
+export class MongoDBAgent implements IAgentMemory {
+  private executeWithRetry: ExecuteWithRetryFn;
+  private getOperationTimeout: GetOperationTimeoutFn;
+
+  constructor(
+    executeWithRetry: ExecuteWithRetryFn,
+    getOperationTimeout: GetOperationTimeoutFn
+  ) {
+    this.executeWithRetry = executeWithRetry;
+    this.getOperationTimeout = getOperationTimeout;
+  }
+
+  public async getAgentPrompt(): Promise<string> {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const metadata = await AgentModel.findOne({
+        id: "agent_metadata"
+      }).maxTimeMS(timeout)
+        .lean<AgentMetadata>();
+      return metadata?.agent_prompt || "";
+    }, "getAgentPrompt()");
+  };
+  
+  public async updateAgentPrompt(prompt: string): Promise<void> {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      await AgentModel.updateOne({
+        id: "agent_metadata",
+      }, { "agent_prompt": prompt }).maxTimeMS(timeout);
+    }, "updateAgentPrompt()");
+  };
+}

package/implements/base.memory.ts

@@ -1,52 +1,264 @@
-import { IMemory } from "node_modules/@ainetwork/adk/dist/esm/modules/memory/base.memory";
+import { IAgentMemory, IIntentMemory, IMemory, IThreadMemory } from "node_modules/@ainetwork/adk/dist/esm/modules/memory/base.memory";
 import mongoose from "mongoose";
 import { loggers } from "@ainetwork/adk/utils/logger";
+import { MongoDBAgent } from "./agent.memory";
+import { MongoDBIntent } from "./intent.memory";
+import { MongoDBThread } from "./thread.memory";
+
+export interface MongoDBMemoryConfig {
+  uri: string;
+  maxReconnectAttempts?: number;
+  reconnectInterval?: number;
+  maxPoolSize?: number;
+  serverSelectionTimeoutMS?: number;
+  socketTimeoutMS?: number;
+  connectTimeoutMS?: number;
+  operationTimeoutMS?: number; // Timeout for database operations
+}
 
 export class MongoDBMemory implements IMemory {
-  private _isConnected: boolean = false;
-  private _uri: string;
+  private static instance: MongoDBMemory;
+  private uri: string;
+  private connected: boolean = false;
+  private reconnectAttempts: number = 0;
+  private maxReconnectAttempts: number;
+  private reconnectInterval: number;
+  private reconnecting: boolean = false;
+  private connectionConfig: mongoose.ConnectOptions;
+  private eventListenersSetup: boolean = false;
+  private operationTimeoutMS: number;
+
+  private agentMemory: MongoDBAgent;
+  private intentMemory: MongoDBIntent;
+  private threadMemory: MongoDBThread;
+
+  constructor(config: string | MongoDBMemoryConfig) {
+    const cfg = typeof config === 'string' ? { uri: config } : config;
+
+    this.uri = cfg.uri;
+    this.maxReconnectAttempts = cfg.maxReconnectAttempts ?? 5;
+    this.reconnectInterval = cfg.reconnectInterval ?? 5000;
+    this.operationTimeoutMS = cfg.operationTimeoutMS ?? 10000; // Default 10 seconds
+    this.connectionConfig = {
+      maxPoolSize: cfg.maxPoolSize ?? 1,
+      serverSelectionTimeoutMS: cfg.serverSelectionTimeoutMS ?? 30000,
+      socketTimeoutMS: cfg.socketTimeoutMS ?? 45000,
+      connectTimeoutMS: cfg.connectTimeoutMS ?? 30000,
+      bufferCommands: false,
+    };
+
+    if (!MongoDBMemory.instance) {
+      MongoDBMemory.instance = this;
+      this.setupMongooseEventListeners();
+    } else {
+      // Use existing instance's connection state
+      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)
+		);
+  }
+
+  public getAgentMemory(): IAgentMemory {
+    return this.agentMemory;
+  }
+
+  public getThreadMemory(): IThreadMemory {
+    return this.threadMemory;
+  }
+
+  public getIntentMemory(): IIntentMemory {
+    return this.intentMemory;
+  }
+
+  private setupMongooseEventListeners(): void {
+    if (this.eventListenersSetup) return;
+
+    this.eventListenersSetup = true;
+
+    mongoose.connection.on("connected", () => {
+      this.connected = true;
+      this.reconnectAttempts = 0;
+      this.reconnecting = false;
+      loggers.agent.info("MongoDB connected successfully");
+    });
+
+    mongoose.connection.on("disconnected", () => {
+      this.connected = false;
+      loggers.agent.warn("MongoDB disconnected");
+      this.handleDisconnection();
+    });
+
+    mongoose.connection.on("error", (error) => {
+      this.connected = false;
+      loggers.agent.error("MongoDB connection error:", error);
+      this.handleDisconnection();
+    });
+
+    mongoose.connection.on("reconnected", () => {
+      this.connected = true;
+      this.reconnectAttempts = 0;
+      this.reconnecting = false;
+      loggers.agent.info("MongoDB reconnected successfully");
+    });
+  }
+
+  private async handleDisconnection(): Promise<void> {
+    if (this.reconnecting) {
+      return;
+    }
+
+    this.reconnecting = true;
+
+    while (this.reconnectAttempts < this.maxReconnectAttempts && !this.isConnected) {
+      this.reconnectAttempts++;
+      loggers.agent.info(
+        `Attempting to reconnect to MongoDB (${this.reconnectAttempts}/${this.maxReconnectAttempts})...`
+      );
+
+      try {
+        await mongoose.connect(this.uri, this.connectionConfig);
+        this.connected = true;
+        this.reconnectAttempts = 0;
+        this.reconnecting = false;
+        loggers.agent.info("MongoDB reconnection successful");
+        return;
+      } catch (error) {
+        loggers.agent.error(
+          `Reconnection attempt ${this.reconnectAttempts} failed:`,
+          error
+        );
+
+        if (this.reconnectAttempts < this.maxReconnectAttempts) {
+          await new Promise((resolve) =>
+            setTimeout(resolve, this.reconnectInterval)
+          );
+        }
+      }
+    }
 
-  constructor(uri: string) {
-    this._uri = uri;
+    this.reconnecting = false;
+
+    if (!this.isConnected) {
+      loggers.agent.error(
+        `Failed to reconnect to MongoDB after ${this.maxReconnectAttempts} attempts`
+      );
+    }
   }
 
   public async connect(): Promise<void> {
-		if (this._isConnected) {
-			return;
-		}
-
-		try {
-      await mongoose.connect(this._uri, {
-        maxPoolSize: 1,
-        serverSelectionTimeoutMS: 30000,
-        socketTimeoutMS: 45000,
-        connectTimeoutMS: 30000,
-        bufferCommands: false,
-      });
-			this._isConnected = true;
-			loggers.agent.info("MongoDB connected successfully");
-		} catch (error) {
-			loggers.agent.error("Failed to connect to MongoDB:", error);
-			throw error;
-		}
+    if (this.connected) {
+      return;
+    }
+
+    try {
+      await mongoose.connect(this.uri, this.connectionConfig);
+      this.connected = true;
+      this.reconnectAttempts = 0;
+    } catch (error) {
+      loggers.agent.error("Failed to connect to MongoDB:", error);
+      throw error;
+    }
   }
 
   public async disconnect(): Promise<void> {
-		if (!this.isConnected) {
-			return;
-		}
+    if (!this.isConnected) {
+      return;
+    }
 
-		try {
-			await mongoose.disconnect();
-			this._isConnected = false;
-			loggers.agent.info("MongoDB disconnected successfully");
-		} catch (error) {
-			loggers.agent.error("Failed to disconnect from MongoDB:", error);
-			throw error;
-		}
+    try {
+      await mongoose.disconnect();
+      this.connected = false;
+    } catch (error) {
+      loggers.agent.error("Failed to disconnect from MongoDB:", error);
+      throw error;
+    }
   }
 
   public isConnected(): boolean {
-    return this._isConnected;
+    return this.connected;
+  }
+
+  private async ensureConnection(): Promise<void> {
+    if (!this.isConnected && !this.reconnecting) {
+      await this.connect();
+    }
+
+    // Wait for reconnection if in progress
+    const maxWaitTime = 30000; // 30 seconds
+    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
+   */
+  protected getOperationTimeout(): number {
+    return this.operationTimeoutMS;
+  }
+
+  /**
+   * Execute a database operation with automatic retry on connection errors
+   * Note: Use mongoose's maxTimeMS option in queries for timeout control
+   */
+  protected async executeWithRetry<T>(
+    operation: () => Promise<T>,
+    operationName: string = "Database operation"
+  ): Promise<T> {
+    await this.ensureConnection();
+
+    try {
+      return await operation();
+    } catch (error: any) {
+      // Check if it's a timeout error from MongoDB
+      if (error.code === 50 || error.message?.includes("operation exceeded time limit")) {
+        loggers.agent.error(`${operationName} exceeded time limit`);
+        throw error;
+      }
+
+      // Check if it's a connection-related error
+      if (
+        error.name === "MongoNetworkError" ||
+        error.name === "MongoServerError" ||
+        error.message?.includes("connection") ||
+        error.message?.includes("disconnect")
+      ) {
+        loggers.agent.warn(
+          `${operationName} failed due to connection issue, attempting reconnection...`
+        );
+
+        await this.ensureConnection();
+
+        // Retry the operation once after reconnection
+        try {
+          return await operation();
+        } catch (retryError: any) {
+          loggers.agent.error(`${operationName} failed after retry:`, retryError);
+          throw retryError;
+        }
+      }
+
+      // If it's not a connection error, just throw it
+      throw error;
+    }
   }
-}
+}

package/implements/intent.memory.ts

@@ -1,35 +1,75 @@
 import type { Intent } from "@ainetwork/adk/types/memory";
 import { IIntentMemory } from "@ainetwork/adk/modules";
-import { MongoDBMemory } from "./base.memory";
 import { IntentModel } from "../models/intent.model";
 
-export class MongoDBIntent extends MongoDBMemory implements IIntentMemory {
+export type ExecuteWithRetryFn = <T>(
+  operation: () => Promise<T>,
+  operationName?: string
+) => Promise<T>;
+
+export type GetOperationTimeoutFn = () => number;
+
+export class MongoDBIntent implements IIntentMemory {
+  private executeWithRetry: ExecuteWithRetryFn;
+  private getOperationTimeout: GetOperationTimeoutFn;
+
+  constructor(
+    executeWithRetry: ExecuteWithRetryFn,
+    getOperationTimeout: GetOperationTimeoutFn
+  ) {
+    this.executeWithRetry = executeWithRetry;
+    this.getOperationTimeout = getOperationTimeout;
+  }
+
   public async getIntent(intentId: string): Promise<Intent | undefined> {
-    const intent = await IntentModel.findOne({ id: intentId }).lean<Intent>();
-    return intent || undefined;
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const intent = await IntentModel.findOne({ id: intentId })
+        .maxTimeMS(timeout)
+        .lean<Intent>();
+      return intent || undefined;
+    }, `getIntent(${intentId})`);
   };
 
-	public async getIntentByName(intentName: string): Promise<Intent | undefined> {
-		const intent = await IntentModel.findOne({ name: intentName }).lean<Intent>();
-    return intent || undefined;
-	}
+  public async getIntentByName(intentName: string): Promise<Intent | undefined> {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const intent = await IntentModel.findOne({ name: intentName })
+        .maxTimeMS(timeout)
+        .lean<Intent>();
+      return intent || undefined;
+    }, `getIntentByName(${intentName})`);
+  }
 
-	public async saveIntent(intent: Intent): Promise<void> {
-    await IntentModel.create(intent);
+  public async saveIntent(intent: Intent): Promise<void> {
+    return this.executeWithRetry(async () => {
+      await IntentModel.create(intent);
+    }, `saveIntent(${intent.id})`);
   };
 
-	public async updateIntent(intentId: string, intent: Intent): Promise<void> {
-    await IntentModel.updateOne({
-      id: intentId,
-    }, intent);
+  public async updateIntent(intentId: string, intent: Intent): Promise<void> {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      await IntentModel.updateOne({
+        id: intentId,
+      }, intent).maxTimeMS(timeout);
+    }, `updateIntent(${intentId})`);
   };
 
-	public async deleteIntent(intentId: string): Promise<void> {
-    await IntentModel.deleteOne({ id: intentId });
+  public async deleteIntent(intentId: string): Promise<void> {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      await IntentModel.deleteOne({ id: intentId }).maxTimeMS(timeout);
+    }, `deleteIntent(${intentId})`);
   };
 
-	public async listIntents(): Promise<Intent[]> {
-    const intents = await IntentModel.find().lean<Intent[]>();
-    return intents;
+  public async listIntents(): Promise<Intent[]> {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const intents = await IntentModel.find()
+        .maxTimeMS(timeout)
+        .lean<Intent[]>();
+      return intents;
+    }, `listIntents()`);
   };
-}
+}