@storecraft/database-mongodb 1.0.12 → 1.0.14

package/vector-store/index.js

@@ -0,0 +1,289 @@
+/**
+ * @import { 
+ *  AIEmbedder, VectorStore 
+ * } from '@storecraft/core/ai/core/types.private.js'
+ * @import {
+ *  Config
+ * } from './types.js'
+ * @import {
+mongo_vectorSearch_pipeline,
+ *  MongoVectorDocument
+ * } from './types.private.js'
+ * 
+ * @import { 
+ *  AnyBulkWriteOperation, Document, AggregationCursor 
+ * } from 'mongodb'
+ * @import { ENV } from '@storecraft/core';
+ */
+
+import { Collection } from 'mongodb';
+import { MongoClient, ServerApiVersion } from 'mongodb';
+
+export const EMBEDDING_KEY_PATH = 'embedding';
+export const NAMESPACE_KEY = 'namespace';
+export const DEFAULT_INDEX_NAME = 'vector_store';
+
+/**
+ * @typedef {VectorStore} Impl
+ */
+
+/**
+ * @description MongoDB Atlas Vector Store
+ * {@link https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-type/#:~:text=You%20can%20use%20the%20vectorSearch,to%20pre%2Dfilter%20your%20data.}
+ * 
+ * @implements {VectorStore}
+ */
+export class MongoVectorStore {
+
+  /** @satisfies {ENV<Config>} */
+  static EnvConfig = /** @type{const} */ ({
+    db_name: 'MONGODB_VECTOR_STORE_DB_NAME',
+    url: 'MONGODB_VECTOR_STORE_URL'
+  });
+
+  /** @type {Config} */
+  config;
+
+  /** @type {MongoClient} */
+  #client
+
+  /**
+   * 
+   * @param {Config} config 
+   */
+  constructor(config) {
+    this.config = {
+      ...config,
+      index_name: config.index_name ?? 'vector_store',
+      similarity: config.similarity ?? 'cosine',
+      dimensions: config.dimensions ?? 1536,
+      options: config.options ?? {
+        ignoreUndefined: true,
+        serverApi: {
+          version: ServerApiVersion.v1,
+          strict: false,
+          deprecationErrors: true,
+        }
+      }
+    };
+  }
+
+  get client() {
+    if(!this.config.db_name || !this.config.url) {
+      throw new Error('MongoVectorStore::client() - missing url or db_name');
+    }
+    
+    this.#client = this.#client ?? new MongoClient(
+      this.config.url, this.config.options
+    );
+    return this.#client;
+  }
+
+  /** @type {VectorStore["onInit"]} */
+  onInit = (app) => {
+    this.config.url ??= app.platform.env[MongoVectorStore.EnvConfig.url ?? 'MONGODB_URL']; 
+    this.config.db_name ??= app.platform.env[MongoVectorStore.EnvConfig.db_name ?? 'MONGODB_DB_NAME'] ?? 'main'; 
+  }
+
+  /** @type {VectorStore["embedder"]} */
+  get embedder() {
+    return this.config.embedder
+  }
+
+  /** @type {Collection<MongoVectorDocument>} */
+  get vector_collection() {
+    return this.client.db(this.config.db_name).collection(this.config.index_name);
+  }
+
+  /** @type {VectorStore["upsertVectors"]} */
+  upsertVectors = async (vectors, documents, options) => {
+    /** @type {MongoVectorDocument[]} */
+    const mongo_docs = documents.map(
+      (doc, ix) => (
+        {
+          updated_at: new Date().toISOString(),
+          embedding: vectors[ix],
+          metadata: doc.metadata,
+          pageContent: doc.pageContent,
+          [NAMESPACE_KEY]: doc.namespace,
+          id: doc.id
+        }
+      )
+    );
+
+    // upsert all docs
+    /** @type {AnyBulkWriteOperation<MongoVectorDocument>[]} */
+    const mongo_replace_ops = mongo_docs.map(
+      (doc) => (
+        {
+          replaceOne: {
+            filter: {
+              id: doc.id
+            },
+            replacement: doc,
+            upsert: true
+          }
+        }
+      )
+    )
+
+    const results = await this.vector_collection.bulkWrite(
+      mongo_replace_ops
+    );
+
+  }
+
+  /** @type {VectorStore["upsertDocuments"]} */
+  upsertDocuments = async (documents, options) => {
+    // first, generate embeddings for the documents
+    const result = await this.embedder.generateEmbeddings(
+      {
+        content: documents.map(
+          doc => (
+            {
+              content: doc.pageContent,
+              type: 'text'
+            }
+          )
+        )
+      }
+    );
+
+    const vectors = result.content;
+
+    return this.upsertVectors(
+      vectors, documents, options
+    )
+  }
+
+  /** @type {VectorStore["delete"]} */
+  delete = async (ids) => {
+    const result = await this.vector_collection.deleteMany(
+      {
+        id: { $in: ids }
+      }
+    );
+    
+  }
+
+  /** @type {VectorStore["similaritySearch"]} */
+  similaritySearch = async (query, k, namespaces) => {
+
+    const embedding_result = await this.embedder.generateEmbeddings(
+      {
+        content: [
+          {
+            content: query, 
+            type: 'text'
+          }
+        ]
+      }
+    );
+    const vector = embedding_result.content[0]
+
+    const agg = [
+      {
+        '$vectorSearch': /** @type {mongo_vectorSearch_pipeline} */ ({
+          index: this.config.index_name,
+          path: EMBEDDING_KEY_PATH,
+          queryVector: vector,
+          numCandidates: k,
+          limit: k,
+          exact: false,
+        })
+      }, {
+        '$project': {
+          '_id': 0,
+          [EMBEDDING_KEY_PATH]: 0,
+          'score': {
+            '$meta': 'vectorSearchScore'
+          }
+        }
+      }
+    ];
+
+    if(Array.isArray(namespaces) && namespaces.length) {
+      agg[0].$vectorSearch.filter = {
+        [NAMESPACE_KEY]: {$in: namespaces}
+      }
+    }
+
+    /** @type {AggregationCursor<MongoVectorDocument>} */
+    const agg_result = this.vector_collection.aggregate(agg);
+    const mongo_vector_docs = await agg_result.toArray();
+
+    return mongo_vector_docs.map(
+      (doc) => {
+        return {
+          score: doc.score,
+          document: {
+            id: doc.id,
+            metadata: /** @type {any}*/ (doc.metadata),
+            pageContent: doc.pageContent,
+            namespace: doc.namespace
+          }
+        }
+      }
+    );
+  }
+
+  /**
+   * @param {boolean} [disconnect_after_finish=true] 
+   * @param {boolean} [delete_index_if_exists_before=false] 
+   * @returns {Promise<boolean>}
+   */
+  createVectorIndex = async (disconnect_after_finish=true, delete_index_if_exists_before=false) => {
+    if(delete_index_if_exists_before) {
+      await this.deleteVectorIndex();
+    }
+
+    const db = this.client.db(this.config.db_name);
+    const collection_name = this.config.index_name;
+    // collection name will have the same name as the index
+    await db.createCollection(collection_name);
+    const index_result = await db.collection(collection_name).createSearchIndex(
+      {
+        name: this.config.index_name,
+        type: 'vectorSearch',
+        definition: {
+          fields: [
+            {
+              type: 'vector',
+              path: EMBEDDING_KEY_PATH,
+              numDimensions: this.config.dimensions,
+              similarity: this.config.similarity
+            },
+            {
+              type: 'filter',
+              path: NAMESPACE_KEY
+            },
+          ]
+        }
+      }
+    );
+
+    if(index_result!==this.config.index_name) {
+      throw new Error('MongoVectorStore::createVectorIndex failed');
+    }
+
+    if(disconnect_after_finish)
+      await this.client.close();
+
+    return true;
+  }
+
+  /**
+   * @returns {Promise<boolean>}
+   */
+  deleteVectorIndex = async () => {
+    const db = this.client.db(this.config.db_name);
+    const collection_name = this.config.index_name;
+    const index_result = await db.collection(collection_name).dropSearchIndex(
+      this.config.index_name
+    );
+
+    return true;
+  }  
+  
+}
+

package/vector-store/types.d.ts

@@ -0,0 +1,44 @@
+
+import type { AIEmbedder } from '@storecraft/core/ai';
+import type { MongoClientOptions } from 'mongodb';
+export * from './index.js';
+
+export type Config = {
+  /**
+   * @description mongo connection url, if absent, will be infered at init
+   * with env `MONGODB_VECTOR_STORE_URL` or `MONGODB_URL`
+   */
+  url?: string;
+
+  /** 
+   * @description the name of the database, if absent, will be infered at init
+   * with env `MONGODB_VECTOR_STORE_NAME` or `MONGODB_NAME` 
+   * @default 'main'
+   */
+  db_name?: string;
+
+  /** 
+   * @description mongo client options 
+   */
+  options?: MongoClientOptions;
+
+  /**
+   * @description The name of the index
+   * @default 'vector_store'
+   */
+  index_name?: string,
+
+  /** 
+   * @description The dimensions of the vectors to be inserted in the index. 
+   * @default 1536
+   */
+  dimensions?: number,
+
+  /**
+   * @description The similiarity metric
+   * @default 'cosine'
+   */
+  similarity?: 'euclidean' | 'cosine' | 'dotProduct',
+
+  embedder: AIEmbedder
+}

package/vector-store/types.private.d.ts

@@ -0,0 +1,44 @@
+import { Document } from "mongodb"
+
+export type MongoVectorDocument = {
+  id: string,
+  metadata: Record<string, any>,
+  embedding: number[],
+  pageContent: string,
+  updated_at: string,
+  score?: number
+  namespace?: string
+}
+
+export type mongo_vectorSearch_pipeline = {
+  /** 
+   * This is required if numCandidates is omitted.
+   * Flag that specifies whether to run ENN or ANN search. Value can be one of the following: 
+   * - false - to run ANN search
+   * - true - to run ENN search 
+   * 
+   * If omitted, defaults to false. 
+   */
+  exact?: boolean
+
+  /**
+   * Any MQL match expression that compares an indexed field with a boolean, date, objectId, number (not decimals), string, or UUID to use as a pre-filter. To learn which query and aggregation pipeline operators Atlas Vector Search supports in your filter, see Atlas Vector Search Pre-Filter.
+   */
+  filter?: Document
+
+  /** Name of the Atlas Vector Search index to use. Atlas Vector Search doesn't return results if you misspell the index name or if the specified index doesn't already exist on the cluster. */
+  index:  string
+
+  /** Number (of type int only) of documents to return in the results. This value can't exceed the value of numCandidates if you specify numCandidates. */
+  limit: number
+
+  /** This field is required if exact is false or omitted. Number of nearest neighbors to use during the search. Value must be less than or equal to (<=) 10000. You can't specify a number less than the number of documents to return (limit). We recommend that you specify a number higher than the number of documents to return (limit) to increase accuracy although this might impact latency. For example, we recommend a ratio of ten to twenty nearest neighbors for a limit of only one document. This overrequest pattern is the recommended way to trade off latency and recall in your ANN searches, and we recommend tuning this on your specific dataset. */
+  numCandidates?: number
+
+  /** Indexed vector type field to search. */
+  path: string
+
+  /** Array of numbers of the BSON double, BSON BinData vector subtype float32, or BSON BinData vector subtype int1 or int8 type that represent the query vector. The number type must match the indexed field value type. Otherwise, Atlas Vector Search doesn't return any results or errors */
+  queryVector: number[]
+
+}