@mastra/mongodb 0.0.2-alpha.0

package/dist/index.js

@@ -0,0 +1,361 @@
+import { MastraVector } from '@mastra/core/vector';
+import { MongoClient } from 'mongodb';
+import { v4 } from 'uuid';
+import { BaseFilterTranslator } from '@mastra/core/vector/filter';
+
+// src/vector/index.ts
+var MongoDBFilterTranslator = class extends BaseFilterTranslator {
+  getSupportedOperators() {
+    return {
+      ...BaseFilterTranslator.DEFAULT_OPERATORS,
+      array: ["$all", "$in", "$nin"],
+      logical: ["$and", "$or", "$not", "$nor"],
+      regex: ["$regex"],
+      custom: ["$size", "$elemMatch"]
+    };
+  }
+  translate(filter) {
+    if (this.isEmpty(filter)) return filter;
+    this.validateFilter(filter);
+    return this.translateNode(filter);
+  }
+  translateNode(node) {
+    if (this.isRegex(node)) {
+      return node;
+    }
+    if (this.isPrimitive(node)) return node;
+    if (Array.isArray(node)) return node;
+    const entries = Object.entries(node);
+    const translatedEntries = entries.map(([key, value]) => {
+      if (this.isOperator(key)) {
+        return [key, this.translateOperatorValue(key, value)];
+      }
+      return [key, this.translateNode(value)];
+    });
+    return Object.fromEntries(translatedEntries);
+  }
+  translateOperatorValue(operator, value) {
+    if (this.isLogicalOperator(operator)) {
+      if (operator === "$not") {
+        if (typeof value !== "object" || value === null || Array.isArray(value)) {
+          throw new Error("$not operator requires an object");
+        }
+        if (this.isEmpty(value)) {
+          throw new Error("$not operator cannot be empty");
+        }
+        return this.translateNode(value);
+      } else {
+        if (!Array.isArray(value)) {
+          throw new Error(`Value for logical operator ${operator} must be an array`);
+        }
+        return value.map((item) => this.translateNode(item));
+      }
+    }
+    if (this.isBasicOperator(operator) || this.isNumericOperator(operator)) {
+      if (value instanceof Date) {
+        return value.toISOString();
+      }
+      return this.normalizeComparisonValue(value);
+    }
+    if (operator === "$elemMatch") {
+      if (typeof value !== "object" || value === null || Array.isArray(value)) {
+        throw new Error(`Value for $elemMatch operator must be an object`);
+      }
+      return this.translateNode(value);
+    }
+    if (this.isArrayOperator(operator)) {
+      if (!Array.isArray(value)) {
+        throw new Error(`Value for array operator ${operator} must be an array`);
+      }
+      return this.normalizeArrayValues(value);
+    }
+    if (this.isElementOperator(operator)) {
+      if (operator === "$exists" && typeof value !== "boolean") {
+        throw new Error(`Value for $exists operator must be a boolean`);
+      }
+      return value;
+    }
+    if (this.isRegexOperator(operator)) {
+      if (!(value instanceof RegExp) && typeof value !== "string") {
+        throw new Error(`Value for ${operator} operator must be a RegExp or string`);
+      }
+      return value;
+    }
+    if (operator === "$size") {
+      if (typeof value !== "number" || !Number.isInteger(value) || value < 0) {
+        throw new Error(`Value for $size operator must be a non-negative integer`);
+      }
+      return value;
+    }
+    throw new Error(`Unsupported operator: ${operator}`);
+  }
+  isEmpty(filter) {
+    return filter === void 0 || filter === null || typeof filter === "object" && Object.keys(filter).length === 0;
+  }
+};
+
+// src/vector/index.ts
+var MongoDBVector = class extends MastraVector {
+  client;
+  db;
+  collections;
+  embeddingFieldName = "embedding";
+  metadataFieldName = "metadata";
+  documentFieldName = "document";
+  collectionForValidation = null;
+  mongoMetricMap = {
+    cosine: "cosine",
+    euclidean: "euclidean",
+    dotproduct: "dotProduct"
+  };
+  constructor({ uri, dbName, options }) {
+    super();
+    this.client = new MongoClient(uri, options);
+    this.db = this.client.db(dbName);
+    this.collections = /* @__PURE__ */ new Map();
+  }
+  // Public methods
+  async connect() {
+    await this.client.connect();
+  }
+  async disconnect() {
+    await this.client.close();
+  }
+  async createIndex(params) {
+    const { indexName, dimension, metric = "cosine" } = params;
+    if (!Number.isInteger(dimension) || dimension <= 0) {
+      throw new Error("Dimension must be a positive integer");
+    }
+    const mongoMetric = this.mongoMetricMap[metric];
+    if (!mongoMetric) {
+      throw new Error(`Invalid metric: "${metric}". Must be one of: cosine, euclidean, dotproduct`);
+    }
+    const collectionExists = await this.db.listCollections({ name: indexName }).hasNext();
+    if (!collectionExists) {
+      await this.db.createCollection(indexName);
+    }
+    const collection = await this.getCollection(indexName);
+    const indexNameInternal = `${indexName}_vector_index`;
+    const embeddingField = this.embeddingFieldName;
+    const numDimensions = dimension;
+    try {
+      await collection.createSearchIndex({
+        definition: {
+          fields: [
+            {
+              type: "vector",
+              path: embeddingField,
+              numDimensions,
+              similarity: mongoMetric
+            }
+          ]
+        },
+        name: indexNameInternal,
+        type: "vectorSearch"
+      });
+    } catch (error) {
+      if (error.codeName !== "IndexAlreadyExists") {
+        throw error;
+      }
+    }
+    await collection.updateOne({ _id: "__index_metadata__" }, { $set: { dimension, metric } }, { upsert: true });
+  }
+  async waitForIndexReady(indexName, timeoutMs = 6e4, checkIntervalMs = 2e3) {
+    const collection = await this.getCollection(indexName, true);
+    const indexNameInternal = `${indexName}_vector_index`;
+    const startTime = Date.now();
+    while (Date.now() - startTime < timeoutMs) {
+      const indexInfo = await collection.listSearchIndexes().toArray();
+      const indexData = indexInfo.find((idx) => idx.name === indexNameInternal);
+      const status = indexData?.status;
+      if (status === "READY") {
+        return;
+      }
+      await new Promise((resolve) => setTimeout(resolve, checkIntervalMs));
+    }
+    throw new Error(`Index "${indexNameInternal}" did not become ready within timeout`);
+  }
+  async upsert(params) {
+    const { indexName, vectors, metadata, ids, documents } = params;
+    const collection = await this.getCollection(indexName);
+    this.collectionForValidation = collection;
+    const stats = await this.describeIndex(indexName);
+    await this.validateVectorDimensions(vectors, stats.dimension);
+    const generatedIds = ids || vectors.map(() => v4());
+    const operations = vectors.map((vector, idx) => {
+      const id = generatedIds[idx];
+      const meta = metadata?.[idx] || {};
+      const doc = documents?.[idx];
+      const normalizedMeta = Object.keys(meta).reduce(
+        (acc, key) => {
+          acc[key] = meta[key] instanceof Date ? meta[key].toISOString() : meta[key];
+          return acc;
+        },
+        {}
+      );
+      const updateDoc = {
+        [this.embeddingFieldName]: vector,
+        [this.metadataFieldName]: normalizedMeta
+      };
+      if (doc !== void 0) {
+        updateDoc[this.documentFieldName] = doc;
+      }
+      return {
+        updateOne: {
+          filter: { _id: id },
+          // '_id' is a string as per MongoDBDocument interface
+          update: { $set: updateDoc },
+          upsert: true
+        }
+      };
+    });
+    await collection.bulkWrite(operations);
+    return generatedIds;
+  }
+  async query(params) {
+    const { indexName, queryVector, topK = 10, filter, includeVector = false, documentFilter } = params;
+    const collection = await this.getCollection(indexName, true);
+    const indexNameInternal = `${indexName}_vector_index`;
+    const mongoFilter = this.transformFilter(filter);
+    const documentMongoFilter = documentFilter ? { [this.documentFieldName]: documentFilter } : {};
+    let combinedFilter = {};
+    if (Object.keys(mongoFilter).length > 0 && Object.keys(documentMongoFilter).length > 0) {
+      combinedFilter = { $and: [mongoFilter, documentMongoFilter] };
+    } else if (Object.keys(mongoFilter).length > 0) {
+      combinedFilter = mongoFilter;
+    } else if (Object.keys(documentMongoFilter).length > 0) {
+      combinedFilter = documentMongoFilter;
+    }
+    const pipeline = [
+      {
+        $vectorSearch: {
+          index: indexNameInternal,
+          queryVector,
+          path: this.embeddingFieldName,
+          numCandidates: 100,
+          limit: topK
+        }
+      },
+      // Apply the filter using $match stage
+      ...Object.keys(combinedFilter).length > 0 ? [{ $match: combinedFilter }] : [],
+      {
+        $set: { score: { $meta: "vectorSearchScore" } }
+      },
+      {
+        $project: {
+          _id: 1,
+          score: 1,
+          metadata: `$${this.metadataFieldName}`,
+          document: `$${this.documentFieldName}`,
+          ...includeVector && { vector: `$${this.embeddingFieldName}` }
+        }
+      }
+    ];
+    try {
+      const results = await collection.aggregate(pipeline).toArray();
+      return results.map((result) => ({
+        id: result._id,
+        score: result.score,
+        metadata: result.metadata,
+        vector: includeVector ? result.vector : void 0,
+        document: result.document
+      }));
+    } catch (error) {
+      console.error("Error during vector search:", error);
+      throw error;
+    }
+  }
+  async listIndexes() {
+    const collections = await this.db.listCollections().toArray();
+    return collections.map((col) => col.name);
+  }
+  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__" });
+    const dimension = metadataDoc?.dimension || 0;
+    const metric = metadataDoc?.metric || "cosine";
+    return {
+      dimension,
+      count,
+      metric
+    };
+  }
+  async deleteIndex(indexName) {
+    const collection = await this.getCollection(indexName, false);
+    if (collection) {
+      await collection.drop();
+      this.collections.delete(indexName);
+    } else {
+      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;
+    }
+    await collection.findOneAndUpdate({ _id: id }, { $set: updateDoc });
+  }
+  async deleteIndexById(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}`);
+    }
+  }
+  // Private methods
+  async getCollection(indexName, throwIfNotExists = true) {
+    if (this.collections.has(indexName)) {
+      return this.collections.get(indexName);
+    }
+    const collection = this.db.collection(indexName);
+    const collectionExists = await this.db.listCollections({ name: indexName }).hasNext();
+    if (!collectionExists && throwIfNotExists) {
+      throw new Error(`Index (Collection) "${indexName}" does not exist`);
+    }
+    this.collections.set(indexName, collection);
+    return collection;
+  }
+  async validateVectorDimensions(vectors, dimension) {
+    if (vectors.length === 0) {
+      throw new Error("No vectors provided for validation");
+    }
+    if (dimension === 0) {
+      dimension = vectors[0] ? vectors[0].length : 0;
+      await this.setIndexDimension(dimension);
+    }
+    for (let i = 0; i < vectors.length; i++) {
+      let v = vectors[i]?.length;
+      if (v !== dimension) {
+        throw new Error(`Vector at index ${i} has invalid dimension ${v}. Expected ${dimension} dimensions.`);
+      }
+    }
+  }
+  async setIndexDimension(dimension) {
+    const collection = this.collectionForValidation;
+    await collection.updateOne({ _id: "__index_metadata__" }, { $set: { dimension } }, { upsert: true });
+  }
+  transformFilter(filter) {
+    const translator = new MongoDBFilterTranslator();
+    if (!filter) return {};
+    return translator.translate(filter);
+  }
+};
+
+export { MongoDBVector };

package/docker-compose.yml

@@ -0,0 +1,8 @@
+services:
+  mongodb:
+    image: mongodb/mongodb-atlas-local
+    environment:
+      MONGODB_INITDB_ROOT_USERNAME: mongodb
+      MONGODB_INITDB_ROOT_PASSWORD: mongodb
+    ports:
+      - 27018:27017

package/eslint.config.js

@@ -0,0 +1,6 @@
+import { createConfig } from '@internal/lint/eslint';
+
+const config = await createConfig();
+
+/** @type {import("eslint").Linter.Config[]} */
+export default [...config];

package/package.json

@@ -0,0 +1,48 @@
+{
+  "name": "@mastra/mongodb",
+  "version": "0.0.2-alpha.0",
+  "description": "MongoDB provider for Mastra - includes vector store capabilities",
+  "type": "module",
+  "main": "dist/index.js",
+  "types": "dist/index.d.ts",
+  "exports": {
+    ".": {
+      "import": {
+        "types": "./dist/index.d.ts",
+        "default": "./dist/index.js"
+      },
+      "require": {
+        "types": "./dist/index.d.cts",
+        "default": "./dist/index.cjs"
+      }
+    },
+    "./package.json": "./package.json"
+  },
+  "dependencies": {
+    "@types/mongodb": "^4.0.7",
+    "cloudflare": "^4.1.0",
+    "mongodb": "^6.15.0",
+    "uuid": "^11.1.0",
+    "@mastra/core": "^0.9.1-alpha.1"
+  },
+  "devDependencies": {
+    "@microsoft/api-extractor": "^7.52.1",
+    "@types/node": "^20.17.27",
+    "eslint": "^9.23.0",
+    "tsup": "^8.4.0",
+    "typescript": "^5.8.2",
+    "vitest": "^3.0.9",
+    "@internal/lint": "0.0.2"
+  },
+  "scripts": {
+    "build": "tsup src/index.ts --format esm,cjs --experimental-dts --clean --treeshake=smallest --splitting",
+    "build:watch": "pnpm build --watch",
+    "pretest": "docker compose up -d && (for i in $(seq 1 30); do docker compose exec -T mongodb mongosh --eval 'db.adminCommand(\"ping\")' --quiet && break || (sleep 1; [ $i -eq 30 ] && exit 1); done)",
+    "test": "vitest run",
+    "posttest": "docker compose down -v",
+    "pretest:watch": "docker compose up -d",
+    "test:watch": "vitest watch",
+    "posttest:watch": "docker compose down -v",
+    "lint": "eslint ."
+  }
+}

package/src/index.ts

@@ -0,0 +1 @@
+export * from './vector';