@mastra/mongodb 0.0.0-vnext-inngest-20250508131921 → 0.0.0-vnextAgentNetwork-20250527091247

package/dist/index.d.ts

@@ -1,7 +1,7 @@
 export { MONGODB_PROMPT } from './_tsup-dts-rollup.js';
-export { MongoDBUpsertArgs } from './_tsup-dts-rollup.js';
-export { MongoDBQueryArgs } from './_tsup-dts-rollup.js';
-export { MongoDBUpsertParams } from './_tsup-dts-rollup.js';
 export { MongoDBUpsertVectorParams } from './_tsup-dts-rollup.js';
 export { MongoDBQueryVectorParams } from './_tsup-dts-rollup.js';
+export { MongoDBIndexReadyParams } from './_tsup-dts-rollup.js';
 export { MongoDBVector } from './_tsup-dts-rollup.js';
+export { MongoDBConfig } from './_tsup-dts-rollup.js';
+export { MongoDBStore } from './_tsup-dts-rollup.js';

package/dist/index.js

@@ -2,6 +2,7 @@ import { MastraVector } from '@mastra/core/vector';
 import { MongoClient } from 'mongodb';
 import { v4 } from 'uuid';
 import { BaseFilterTranslator } from '@mastra/core/vector/filter';
+import { MastraStorage, TABLE_THREADS, TABLE_MESSAGES, TABLE_TRACES, TABLE_WORKFLOW_SNAPSHOT, TABLE_EVALS } from '@mastra/core/storage';
 
 // src/vector/index.ts
 var MongoDBFilterTranslator = class extends BaseFilterTranslator {
@@ -119,8 +120,7 @@ var MongoDBVector = class extends MastraVector {
   async disconnect() {
     await this.client.close();
   }
-  async createIndex(params) {
-    const { indexName, dimension, metric = "cosine" } = params;
+  async createIndex({ indexName, dimension, metric = "cosine" }) {
     if (!Number.isInteger(dimension) || dimension <= 0) {
       throw new Error("Dimension must be a positive integer");
     }
@@ -158,7 +158,19 @@ var MongoDBVector = class extends MastraVector {
     }
     await collection.updateOne({ _id: "__index_metadata__" }, { $set: { dimension, metric } }, { upsert: true });
   }
-  async waitForIndexReady(indexName, timeoutMs = 6e4, checkIntervalMs = 2e3) {
+  /**
+   * Waits for the index to be ready.
+   *
+   * @param {string} indexName - The name of the index to wait for
+   * @param {number} timeoutMs - The maximum time in milliseconds to wait for the index to be ready (default: 60000)
+   * @param {number} checkIntervalMs - The interval in milliseconds at which to check if the index is ready (default: 2000)
+   * @returns A promise that resolves when the index is ready
+   */
+  async waitForIndexReady({
+    indexName,
+    timeoutMs = 6e4,
+    checkIntervalMs = 2e3
+  }) {
     const collection = await this.getCollection(indexName, true);
     const indexNameInternal = `${indexName}_vector_index`;
     const startTime = Date.now();
@@ -173,11 +185,10 @@ var MongoDBVector = class extends MastraVector {
     }
     throw new Error(`Index "${indexNameInternal}" did not become ready within timeout`);
   }
-  async upsert(params) {
-    const { indexName, vectors, metadata, ids, documents } = params;
+  async upsert({ indexName, vectors, metadata, ids, documents }) {
     const collection = await this.getCollection(indexName);
     this.collectionForValidation = collection;
-    const stats = await this.describeIndex(indexName);
+    const stats = await this.describeIndex({ indexName });
     await this.validateVectorDimensions(vectors, stats.dimension);
     const generatedIds = ids || vectors.map(() => v4());
     const operations = vectors.map((vector, idx) => {
@@ -210,8 +221,14 @@ var MongoDBVector = class extends MastraVector {
     await collection.bulkWrite(operations);
     return generatedIds;
   }
-  async query(params) {
-    const { indexName, queryVector, topK = 10, filter, includeVector = false, documentFilter } = params;
+  async query({
+    indexName,
+    queryVector,
+    topK = 10,
+    filter,
+    includeVector = false,
+    documentFilter
+  }) {
     const collection = await this.getCollection(indexName, true);
     const indexNameInternal = `${indexName}_vector_index`;
     const mongoFilter = this.transformFilter(filter);
@@ -267,7 +284,13 @@ var MongoDBVector = class extends MastraVector {
     const collections = await this.db.listCollections().toArray();
     return collections.map((col) => col.name);
   }
-  async describeIndex(indexName) {
+  /**
+   * Retrieves statistics about a vector index.
+   *
+   * @param {string} indexName - The name of the index to describe
+   * @returns A promise that resolves to the index statistics including dimension, count and metric
+   */
+  async describeIndex({ indexName }) {
     const collection = await this.getCollection(indexName, true);
     const count = await collection.countDocuments({ _id: { $ne: "__index_metadata__" } });
     const metadataDoc = await collection.findOne({ _id: "__index_metadata__" });
@@ -279,7 +302,7 @@ var MongoDBVector = class extends MastraVector {
       metric
     };
   }
-  async deleteIndex(indexName) {
+  async deleteIndex({ indexName }) {
     const collection = await this.getCollection(indexName, false);
     if (collection) {
       await collection.drop();
@@ -288,33 +311,56 @@ var MongoDBVector = class extends MastraVector {
       throw new Error(`Index (Collection) "${indexName}" does not exist`);
     }
   }
-  async updateIndexById(indexName, id, update) {
-    if (!update.vector && !update.metadata) {
-      throw new Error("No updates provided");
-    }
-    const collection = await this.getCollection(indexName, true);
-    const updateDoc = {};
-    if (update.vector) {
-      updateDoc[this.embeddingFieldName] = update.vector;
-    }
-    if (update.metadata) {
-      const normalizedMeta = Object.keys(update.metadata).reduce(
-        (acc, key) => {
-          acc[key] = update.metadata[key] instanceof Date ? update.metadata[key].toISOString() : update.metadata[key];
-          return acc;
-        },
-        {}
-      );
-      updateDoc[this.metadataFieldName] = normalizedMeta;
+  /**
+   * Updates a vector by its ID with the provided vector and/or metadata.
+   * @param indexName - The name of the index containing the vector.
+   * @param id - The ID of the vector to update.
+   * @param update - An object containing the vector and/or metadata to update.
+   * @param update.vector - An optional array of numbers representing the new vector.
+   * @param update.metadata - An optional record containing the new metadata.
+   * @returns A promise that resolves when the update is complete.
+   * @throws Will throw an error if no updates are provided or if the update operation fails.
+   */
+  async updateVector({ indexName, id, update }) {
+    try {
+      if (!update.vector && !update.metadata) {
+        throw new Error("No updates provided");
+      }
+      const collection = await this.getCollection(indexName, true);
+      const updateDoc = {};
+      if (update.vector) {
+        const stats = await this.describeIndex({ indexName });
+        await this.validateVectorDimensions([update.vector], stats.dimension);
+        updateDoc[this.embeddingFieldName] = update.vector;
+      }
+      if (update.metadata) {
+        const normalizedMeta = Object.keys(update.metadata).reduce(
+          (acc, key) => {
+            acc[key] = update.metadata[key] instanceof Date ? update.metadata[key].toISOString() : update.metadata[key];
+            return acc;
+          },
+          {}
+        );
+        updateDoc[this.metadataFieldName] = normalizedMeta;
+      }
+      await collection.findOneAndUpdate({ _id: id }, { $set: updateDoc });
+    } catch (error) {
+      throw new Error(`Failed to update vector by id: ${id} for index name: ${indexName}: ${error.message}`);
     }
-    await collection.findOneAndUpdate({ _id: id }, { $set: updateDoc });
   }
-  async deleteIndexById(indexName, id) {
+  /**
+   * Deletes a vector by its ID.
+   * @param indexName - The name of the index containing the vector.
+   * @param id - The ID of the vector to delete.
+   * @returns A promise that resolves when the deletion is complete.
+   * @throws Will throw an error if the deletion operation fails.
+   */
+  async deleteVector({ indexName, id }) {
     try {
       const collection = await this.getCollection(indexName, true);
       await collection.deleteOne({ _id: id });
     } catch (error) {
-      throw new Error(`Failed to delete index by id: ${id} for index name: ${indexName}: ${error.message}`);
+      throw new Error(`Failed to delete vector by id: ${id} for index name: ${indexName}: ${error.message}`);
     }
   }
   // Private methods
@@ -355,6 +401,524 @@ var MongoDBVector = class extends MastraVector {
     return translator.translate(filter);
   }
 };
+function safelyParseJSON(jsonString) {
+  try {
+    return JSON.parse(jsonString);
+  } catch {
+    return {};
+  }
+}
+var MongoDBStore = class extends MastraStorage {
+  #isConnected = false;
+  #client;
+  #db;
+  #dbName;
+  constructor(config) {
+    super({ name: "MongoDBStore" });
+    this.#isConnected = false;
+    if (!config.url?.trim().length) {
+      throw new Error(
+        "MongoDBStore: url must be provided and cannot be empty. Passing an empty string may cause fallback to local MongoDB defaults."
+      );
+    }
+    if (!config.dbName?.trim().length) {
+      throw new Error(
+        "MongoDBStore: dbName must be provided and cannot be empty. Passing an empty string may cause fallback to local MongoDB defaults."
+      );
+    }
+    this.#dbName = config.dbName;
+    this.#client = new MongoClient(config.url, config.options);
+  }
+  async getConnection() {
+    if (this.#isConnected) {
+      return this.#db;
+    }
+    await this.#client.connect();
+    this.#db = this.#client.db(this.#dbName);
+    this.#isConnected = true;
+    return this.#db;
+  }
+  async getCollection(collectionName) {
+    const db = await this.getConnection();
+    return db.collection(collectionName);
+  }
+  async createTable() {
+  }
+  async clearTable({ tableName }) {
+    try {
+      const collection = await this.getCollection(tableName);
+      await collection.deleteMany({});
+    } catch (error) {
+      if (error instanceof Error) {
+        this.logger.error(error.message);
+      }
+    }
+  }
+  async insert({ tableName, record }) {
+    try {
+      const collection = await this.getCollection(tableName);
+      await collection.insertOne(record);
+    } catch (error) {
+      this.logger.error(`Error upserting into table ${tableName}: ${error}`);
+      throw error;
+    }
+  }
+  async batchInsert({ tableName, records }) {
+    if (!records.length) {
+      return;
+    }
+    try {
+      const collection = await this.getCollection(tableName);
+      await collection.insertMany(records);
+    } catch (error) {
+      this.logger.error(`Error upserting into table ${tableName}: ${error}`);
+      throw error;
+    }
+  }
+  async load({ tableName, keys }) {
+    this.logger.info(`Loading ${tableName} with keys ${JSON.stringify(keys)}`);
+    try {
+      const collection = await this.getCollection(tableName);
+      return await collection.find(keys).toArray();
+    } catch (error) {
+      this.logger.error(`Error loading ${tableName} with keys ${JSON.stringify(keys)}: ${error}`);
+      throw error;
+    }
+  }
+  async getThreadById({ threadId }) {
+    try {
+      const collection = await this.getCollection(TABLE_THREADS);
+      const result = await collection.findOne({ id: threadId });
+      if (!result) {
+        return null;
+      }
+      return {
+        ...result,
+        metadata: typeof result.metadata === "string" ? JSON.parse(result.metadata) : result.metadata
+      };
+    } catch (error) {
+      this.logger.error(`Error loading thread with ID ${threadId}: ${error}`);
+      throw error;
+    }
+  }
+  async getThreadsByResourceId({ resourceId }) {
+    try {
+      const collection = await this.getCollection(TABLE_THREADS);
+      const results = await collection.find({ resourceId }).toArray();
+      if (!results.length) {
+        return [];
+      }
+      return results.map((result) => ({
+        ...result,
+        metadata: typeof result.metadata === "string" ? JSON.parse(result.metadata) : result.metadata
+      }));
+    } catch (error) {
+      this.logger.error(`Error loading threads by resourceId ${resourceId}: ${error}`);
+      throw error;
+    }
+  }
+  async saveThread({ thread }) {
+    try {
+      const collection = await this.getCollection(TABLE_THREADS);
+      await collection.updateOne(
+        { id: thread.id },
+        {
+          $set: {
+            ...thread,
+            metadata: JSON.stringify(thread.metadata)
+          }
+        },
+        { upsert: true }
+      );
+      return thread;
+    } catch (error) {
+      this.logger.error(`Error saving thread ${thread.id}: ${error}`);
+      throw error;
+    }
+  }
+  async updateThread({
+    id,
+    title,
+    metadata
+  }) {
+    const thread = await this.getThreadById({ threadId: id });
+    if (!thread) {
+      throw new Error(`Thread ${id} not found`);
+    }
+    const updatedThread = {
+      ...thread,
+      title,
+      metadata: {
+        ...thread.metadata,
+        ...metadata
+      }
+    };
+    try {
+      const collection = await this.getCollection(TABLE_THREADS);
+      await collection.updateOne(
+        { id },
+        {
+          $set: {
+            title,
+            metadata: JSON.stringify(updatedThread.metadata)
+          }
+        }
+      );
+    } catch (error) {
+      this.logger.error(`Error updating thread ${id}:) ${error}`);
+      throw error;
+    }
+    return updatedThread;
+  }
+  async deleteThread({ threadId }) {
+    try {
+      const collectionMessages = await this.getCollection(TABLE_MESSAGES);
+      await collectionMessages.deleteMany({ thread_id: threadId });
+      const collectionThreads = await this.getCollection(TABLE_THREADS);
+      await collectionThreads.deleteOne({ id: threadId });
+    } catch (error) {
+      this.logger.error(`Error deleting thread ${threadId}: ${error}`);
+      throw error;
+    }
+  }
+  async getMessages({ threadId, selectBy }) {
+    try {
+      const limit = typeof selectBy?.last === "number" ? selectBy.last : 40;
+      const include = selectBy?.include || [];
+      let messages = [];
+      let allMessages = [];
+      const collection = await this.getCollection(TABLE_MESSAGES);
+      allMessages = (await collection.find({ thread_id: threadId }).sort({ createdAt: -1 }).toArray()).map(
+        (row) => this.parseRow(row)
+      );
+      if (include.length) {
+        const idToIndex = /* @__PURE__ */ new Map();
+        allMessages.forEach((msg, idx) => {
+          idToIndex.set(msg.id, idx);
+        });
+        const selectedIndexes = /* @__PURE__ */ new Set();
+        for (const inc of include) {
+          const idx = idToIndex.get(inc.id);
+          if (idx === void 0) continue;
+          for (let i = 1; i <= (inc.withPreviousMessages || 0); i++) {
+            if (idx + i < allMessages.length) selectedIndexes.add(idx + i);
+          }
+          selectedIndexes.add(idx);
+          for (let i = 1; i <= (inc.withNextMessages || 0); i++) {
+            if (idx - i >= 0) selectedIndexes.add(idx - i);
+          }
+        }
+        messages.push(
+          ...Array.from(selectedIndexes).map((i) => allMessages[i]).filter((m) => !!m)
+        );
+      }
+      const excludeIds = new Set(messages.map((m) => m.id));
+      for (const msg of allMessages) {
+        if (messages.length >= limit) break;
+        if (!excludeIds.has(msg.id)) {
+          messages.push(msg);
+        }
+      }
+      messages.sort((a, b) => a.createdAt.getTime() - b.createdAt.getTime());
+      return messages.slice(0, limit);
+    } catch (error) {
+      this.logger.error("Error getting messages:", error);
+      throw error;
+    }
+  }
+  async saveMessages({ messages }) {
+    if (!messages.length) {
+      return messages;
+    }
+    const threadId = messages[0]?.threadId;
+    if (!threadId) {
+      this.logger.error("Thread ID is required to save messages");
+      throw new Error("Thread ID is required");
+    }
+    try {
+      const messagesToInsert = messages.map((message) => {
+        const time = message.createdAt || /* @__PURE__ */ new Date();
+        return {
+          id: message.id,
+          thread_id: threadId,
+          content: typeof message.content === "string" ? message.content : JSON.stringify(message.content),
+          role: message.role,
+          type: message.type,
+          resourceId: message.resourceId,
+          createdAt: time instanceof Date ? time.toISOString() : time
+        };
+      });
+      const collection = await this.getCollection(TABLE_MESSAGES);
+      await collection.insertMany(messagesToInsert);
+      return messages;
+    } catch (error) {
+      this.logger.error("Failed to save messages in database: " + error?.message);
+      throw error;
+    }
+  }
+  async getTraces({
+    name,
+    scope,
+    page,
+    perPage,
+    attributes,
+    filters
+  } = {
+    page: 0,
+    perPage: 100
+  }) {
+    const limit = perPage;
+    const offset = page * perPage;
+    const query = {};
+    if (name) {
+      query["name"] = `%${name}%`;
+    }
+    if (scope) {
+      query["scope"] = scope;
+    }
+    if (attributes) {
+      Object.keys(attributes).forEach((key) => {
+        query[`attributes.${key}`] = attributes[key];
+      });
+    }
+    if (filters) {
+      Object.entries(filters).forEach(([key, value]) => {
+        query[key] = value;
+      });
+    }
+    const collection = await this.getCollection(TABLE_TRACES);
+    const result = await collection.find(query, {
+      sort: { startTime: -1 }
+    }).limit(limit).skip(offset).toArray();
+    return result.map((row) => ({
+      id: row.id,
+      parentSpanId: row.parentSpanId,
+      traceId: row.traceId,
+      name: row.name,
+      scope: row.scope,
+      kind: row.kind,
+      status: safelyParseJSON(row.status),
+      events: safelyParseJSON(row.events),
+      links: safelyParseJSON(row.links),
+      attributes: safelyParseJSON(row.attributes),
+      startTime: row.startTime,
+      endTime: row.endTime,
+      other: safelyParseJSON(row.other),
+      createdAt: row.createdAt
+    }));
+  }
+  async getWorkflowRuns({
+    workflowName,
+    fromDate,
+    toDate,
+    limit,
+    offset
+  } = {}) {
+    const query = {};
+    if (workflowName) {
+      query["workflow_name"] = workflowName;
+    }
+    if (fromDate || toDate) {
+      query["createdAt"] = {};
+      if (fromDate) {
+        query["createdAt"]["$gte"] = fromDate;
+      }
+      if (toDate) {
+        query["createdAt"]["$lte"] = toDate;
+      }
+    }
+    const collection = await this.getCollection(TABLE_WORKFLOW_SNAPSHOT);
+    let total = 0;
+    if (limit !== void 0 && offset !== void 0) {
+      total = await collection.countDocuments(query);
+    }
+    const request = collection.find(query).sort({ createdAt: "desc" });
+    if (limit) {
+      request.limit(limit);
+    }
+    if (offset) {
+      request.skip(offset);
+    }
+    const result = await request.toArray();
+    const runs = result.map((row) => {
+      let parsedSnapshot = row.snapshot;
+      if (typeof parsedSnapshot === "string") {
+        try {
+          parsedSnapshot = JSON.parse(row.snapshot);
+        } catch (e) {
+          console.warn(`Failed to parse snapshot for workflow ${row.workflow_name}: ${e}`);
+        }
+      }
+      return {
+        workflowName: row.workflow_name,
+        runId: row.run_id,
+        snapshot: parsedSnapshot,
+        createdAt: new Date(row.createdAt),
+        updatedAt: new Date(row.updatedAt)
+      };
+    });
+    return { runs, total: total || runs.length };
+  }
+  async getEvalsByAgentName(agentName, type) {
+    try {
+      const query = {
+        agent_name: agentName
+      };
+      if (type === "test") {
+        query["test_info"] = { $ne: null };
+      }
+      if (type === "live") {
+        query["test_info"] = null;
+      }
+      const collection = await this.getCollection(TABLE_EVALS);
+      const documents = await collection.find(query).sort({ created_at: "desc" }).toArray();
+      const result = documents.map((row) => this.transformEvalRow(row));
+      return result.filter((row) => {
+        if (type === "live") {
+          return !Boolean(row.testInfo?.testPath);
+        }
+        if (type === "test") {
+          return row.testInfo?.testPath !== null;
+        }
+        return true;
+      });
+    } catch (error) {
+      if (error instanceof Error && error.message.includes("no such table")) {
+        return [];
+      }
+      this.logger.error("Failed to get evals for the specified agent: " + error?.message);
+      throw error;
+    }
+  }
+  async persistWorkflowSnapshot({
+    workflowName,
+    runId,
+    snapshot
+  }) {
+    try {
+      const now = (/* @__PURE__ */ new Date()).toISOString();
+      const collection = await this.getCollection(TABLE_WORKFLOW_SNAPSHOT);
+      await collection.updateOne(
+        { workflow_name: workflowName, run_id: runId },
+        {
+          $set: {
+            snapshot: JSON.stringify(snapshot),
+            updatedAt: now
+          },
+          $setOnInsert: {
+            createdAt: now
+          }
+        },
+        { upsert: true }
+      );
+    } catch (error) {
+      this.logger.error(`Error persisting workflow snapshot: ${error}`);
+      throw error;
+    }
+  }
+  async loadWorkflowSnapshot({
+    workflowName,
+    runId
+  }) {
+    try {
+      const result = await this.load({
+        tableName: TABLE_WORKFLOW_SNAPSHOT,
+        keys: {
+          workflow_name: workflowName,
+          run_id: runId
+        }
+      });
+      if (!result?.length) {
+        return null;
+      }
+      return JSON.parse(result[0].snapshot);
+    } catch (error) {
+      console.error("Error loading workflow snapshot:", error);
+      throw error;
+    }
+  }
+  async getWorkflowRunById({
+    runId,
+    workflowName
+  }) {
+    try {
+      const query = {};
+      if (runId) {
+        query["run_id"] = runId;
+      }
+      if (workflowName) {
+        query["workflow_name"] = workflowName;
+      }
+      const collection = await this.getCollection(TABLE_WORKFLOW_SNAPSHOT);
+      const result = await collection.findOne(query);
+      if (!result) {
+        return null;
+      }
+      return this.parseWorkflowRun(result);
+    } catch (error) {
+      console.error("Error getting workflow run by ID:", error);
+      throw error;
+    }
+  }
+  parseWorkflowRun(row) {
+    let parsedSnapshot = row.snapshot;
+    if (typeof parsedSnapshot === "string") {
+      try {
+        parsedSnapshot = JSON.parse(row.snapshot);
+      } catch (e) {
+        console.warn(`Failed to parse snapshot for workflow ${row.workflow_name}: ${e}`);
+      }
+    }
+    return {
+      workflowName: row.workflow_name,
+      runId: row.run_id,
+      snapshot: parsedSnapshot,
+      createdAt: row.createdAt,
+      updatedAt: row.updatedAt,
+      resourceId: row.resourceId
+    };
+  }
+  parseRow(row) {
+    let content = row.content;
+    try {
+      content = JSON.parse(row.content);
+    } catch {
+    }
+    return {
+      id: row.id,
+      content,
+      role: row.role,
+      type: row.type,
+      createdAt: new Date(row.createdAt),
+      threadId: row.thread_id
+    };
+  }
+  transformEvalRow(row) {
+    let testInfoValue = null;
+    if (row.test_info) {
+      try {
+        testInfoValue = typeof row.test_info === "string" ? JSON.parse(row.test_info) : row.test_info;
+      } catch (e) {
+        console.warn("Failed to parse test_info:", e);
+      }
+    }
+    return {
+      input: row.input,
+      output: row.output,
+      result: row.result,
+      agentName: row.agent_name,
+      metricName: row.metric_name,
+      instructions: row.instructions,
+      testInfo: testInfoValue,
+      globalRunId: row.global_run_id,
+      runId: row.run_id,
+      createdAt: row.created_at
+    };
+  }
+  async close() {
+    await this.#client.close();
+  }
+};
 
 // src/vector/prompt.ts
 var MONGODB_PROMPT = `When querying MongoDB Vector, you can ONLY use the operators listed below. Any other operators will be rejected.
@@ -451,4 +1015,4 @@ Example Complex Query:
   ]
 }`;
 
-export { MONGODB_PROMPT, MongoDBVector };
+export { MONGODB_PROMPT, MongoDBStore, MongoDBVector };