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

package/implements/thread.memory.ts

@@ -1,115 +1,141 @@
 import type { MessageObject, ThreadMetadata, ThreadObject, ThreadType } from "@ainetwork/adk/types/memory";
 import { MessageRole } from "@ainetwork/adk/types/memory";
 import { IThreadMemory } from "@ainetwork/adk/modules";
-import { MongoDBMemory } from "./base.memory";
 import { ThreadDocument, ThreadModel } from "../models/threads.model";
 import { MessageDocument, MessageModel } from "../models/messages.model";
 import { loggers } from "@ainetwork/adk/utils/logger";
 
-export class MongoDBThread extends MongoDBMemory implements IThreadMemory {
-  constructor(uri: string) {
-    super(uri);
+export type ExecuteWithRetryFn = <T>(
+  operation: () => Promise<T>,
+  operationName?: string
+) => Promise<T>;
+
+export type GetOperationTimeoutFn = () => number;
+
+export class MongoDBThread implements IThreadMemory {
+  private executeWithRetry: ExecuteWithRetryFn;
+  private getOperationTimeout: GetOperationTimeoutFn;
+
+  constructor(
+    executeWithRetry: ExecuteWithRetryFn,
+    getOperationTimeout: GetOperationTimeoutFn
+  ) {
+    this.executeWithRetry = executeWithRetry;
+    this.getOperationTimeout = getOperationTimeout;
   }
 
   public async getThread(
     userId: string,
     threadId: string
   ): Promise<ThreadObject | undefined> {
-    const thread = await ThreadModel.findOne({ threadId, userId });
-		const messages = await MessageModel.find({ threadId, userId }).sort({
-			timestamp: 1,
-		});
-
-    if (!thread) return undefined;
-
-		loggers.agent.debug(`Found ${messages.length} messages for thread ${threadId}`);
-
-		const threadObject: ThreadObject = { 
-      threadId: thread.threadId, 
-      userId: thread.userId,
-      type: thread.type as ThreadType,
-      title: thread.title || "New thread",
-      messages: []
-    };
-		messages.forEach((message: MessageDocument) => {
-			threadObject.messages.push({
-        messageId: message.messageId,
-				role: message.role as MessageRole,
-				content: message.content,
-				timestamp: message.timestamp,
-				metadata: message.metadata,
-			});
-		});
-
-		return threadObject;
+    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 undefined;
+
+      loggers.agent.debug(`Found ${messages.length} messages for thread ${threadId}`);
+
+      const threadObject: ThreadObject = {
+        threadId: thread.threadId,
+        userId: thread.userId,
+        type: thread.type as ThreadType,
+        title: thread.title || "New thread",
+        messages: []
+      };
+      messages.forEach((message: MessageDocument) => {
+        threadObject.messages.push({
+          messageId: message.messageId,
+          role: message.role as MessageRole,
+          content: message.content,
+          timestamp: message.timestamp,
+          metadata: message.metadata,
+        });
+      });
+
+      return threadObject;
+    }, `getThread(${userId}, ${threadId})`);
   };
 
-	public async createThread(
-		type: ThreadType,
-		userId: string,
-		threadId: string,
-		title: string,
+  public async createThread(
+    type: ThreadType,
+    userId: string,
+    threadId: string,
+    title: string,
   ): Promise<ThreadObject> {
-    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})`);
   };
 
-	public async addMessagesToThread(
+  public async addMessagesToThread(
     userId: string,
     threadId: string,
     messages: MessageObject[]
   ): Promise<void> {
-    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})`);
   };
 
-	public async deleteThread(userId: string, threadId: string): Promise<void> {
-		const messages = await MessageModel.find({ userId, threadId }).sort({
-			timestamp: 1,
-		});
-
-		messages?.forEach((message: MessageDocument) => {
-      message.deleteOne();
-		});
-    
-    const thread = await ThreadModel.findOne({ userId, threadId });
-    thread?.deleteOne();
+  public async deleteThread(userId: string, threadId: string): Promise<void> {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const messages = await MessageModel.find({ userId, threadId })
+        .sort({ timestamp: 1 })
+        .maxTimeMS(timeout);
+
+      messages?.forEach((message: MessageDocument) => {
+        message.deleteOne();
+      });
+
+      const thread = await ThreadModel.findOne({ userId, threadId }).maxTimeMS(timeout);
+      thread?.deleteOne();
+    }, `deleteThread(${userId}, ${threadId})`);
   };
 
-	public async listThreads(userId: string): Promise<ThreadMetadata[]> {
-    const threads = await ThreadModel.find({ userId }).sort({
-      updated_at: -1,
-    });
-    const data: ThreadMetadata[] = threads.map((thread: ThreadDocument) => {
-      return {
-        type: thread.type,
-        userId,
-        threadId: thread.threadId,
-        title: thread.title,
-        updatedAt: thread.updated_at
-      } as ThreadMetadata;
-    })
-    return data;
+  public async listThreads(userId: string): Promise<ThreadMetadata[]> {
+    return this.executeWithRetry(async () => {
+      const timeout = this.getOperationTimeout();
+      const threads = await ThreadModel.find({ userId })
+        .sort({ updated_at: -1 })
+        .maxTimeMS(timeout);
+      const data: ThreadMetadata[] = threads.map((thread: ThreadDocument) => {
+        return {
+          type: thread.type,
+          userId,
+          threadId: thread.threadId,
+          title: thread.title,
+          updatedAt: thread.updated_at
+        } as ThreadMetadata;
+      })
+      return data;
+    }, `listThreads(${userId})`);
   };
-}
+}

package/index.ts

@@ -1,2 +1,2 @@
-export { MongoDBThread } from "./implements/thread.memory";
-export { MongoDBIntent } from "./implements/intent.memory";
+export { MongoDBMemory } from "./implements/base.memory";
+export type { MongoDBMemoryConfig } from "./implements/base.memory";

package/models/agent.model.ts

@@ -0,0 +1,16 @@
+import { type Document, Schema } from "mongoose";
+import mongoose from "mongoose";
+
+export const AgentObjectSchema = new Schema(
+	{
+		agent_prompt: {
+			type: String,
+		},
+	},
+);
+
+export interface AgentDocument extends Document {
+	agent_prompt: string;
+}
+
+export const AgentModel = mongoose.model<AgentDocument>("Agent", AgentObjectSchema);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ainetwork/adk-provider-memory-mongodb",
-  "version": "0.2.3",
+  "version": "0.3.0",
   "author": "AI Network (https://ainetwork.ai)",
   "type": "module",
   "engines": {
@@ -28,7 +28,7 @@
     "clean": "rm -rf dist"
   },
   "dependencies": {
-    "@ainetwork/adk": "0.2.9",
+    "@ainetwork/adk": "^0.3.0",
     "mongoose": "^8.16.5"
   },
   "devDependencies": {
@@ -38,5 +38,5 @@
   "publishConfig": {
     "access": "public"
   },
-  "gitHead": "da5af7747935eeb8cedeb932cb2e918da6dc36c3"
+  "gitHead": "77a4f8c8bb228f45383eea93f153c21bac3f8b3b"
 }