feat: integrate with azure cosmos db mongo vCore (#1500)

Co-authored-by: root <root@CPC-yangq-FRSGK>
Co-authored-by: Alex Yang <himself65@outlook.com>

.changeset/tough-rings-roll.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+feat: add Azure Cosmos DB mongo vCore DocumentStore, IndexStore, KVStore

packages/llamaindex/src/storage/docStore/AzureCosmosMongovCoreDocumentStore.ts

@@ -0,0 +1,49 @@
+import { MongoClient } from "mongodb";
+import { AzureCosmosVCoreKVStore } from "../kvStore/AzureCosmosMongovCoreKVStore.js";
+import { KVDocumentStore } from "./KVDocumentStore.js";
+
+const DEFAULT_DATABASE = "DocumentStoreDB";
+const DEFAULT_COLLECTION = "DocumentStoreCollection";
+
+export interface AzureCosmosVCoreDocumentStoreArgs {
+  azureCosmosVCoreKVStore: AzureCosmosVCoreKVStore;
+  namespace?: string;
+}
+
+export class AzureCosmosVCoreDocumentStore extends KVDocumentStore {
+  constructor({
+    azureCosmosVCoreKVStore,
+    namespace,
+  }: AzureCosmosVCoreDocumentStoreArgs) {
+    super(azureCosmosVCoreKVStore, namespace);
+  }
+
+  /**
+   * Static method for creating an instance using a MongoClient.
+   * @returns Instance of AzureCosmosVCoreDocumentStore
+   * @param mongoClient - MongoClient instance
+   * @param dbName - Database name
+   * @param collectionName - Collection name
+   * @example
+   * ```ts
+   * const mongoClient = new MongoClient("mongodb://localhost:27017");
+   * const indexStore = AzureCosmosVCoreDocumentStore.fromMongoClient(mongoClient, "my_db", "my_collection");
+   * ```
+   */
+  static fromMongoClient(
+    mongoClient: MongoClient,
+    dbName: string = DEFAULT_DATABASE,
+    collectionName: string = DEFAULT_COLLECTION,
+  ) {
+    const azureCosmosVCoreKVStore = new AzureCosmosVCoreKVStore({
+      mongoClient,
+      dbName,
+      collectionName,
+    });
+    const namespace = `${dbName}.${collectionName}`;
+    return new AzureCosmosVCoreDocumentStore({
+      azureCosmosVCoreKVStore,
+      namespace,
+    });
+  }
+}

packages/llamaindex/src/storage/index.ts

@@ -3,12 +3,15 @@ export * from "@llamaindex/core/storage/doc-store";
 export * from "@llamaindex/core/storage/index-store";
 export * from "@llamaindex/core/storage/kv-store";
 export * from "./chatStore/AzureCosmosNoSqlChatStore.js";
+export * from "./docStore/AzureCosmosMongovCoreDocumentStore.js";
 export * from "./docStore/AzureCosmosNoSqlDocumentStore.js";
 export { PostgresDocumentStore } from "./docStore/PostgresDocumentStore.js";
 export { SimpleDocumentStore } from "./docStore/SimpleDocumentStore.js";
 export * from "./FileSystem.js";
+export * from "./indexStore/AzureCosmosMongovCoreIndexStore.js";
 export * from "./indexStore/AzureCosmosNoSqlIndexStore.js";
 export { PostgresIndexStore } from "./indexStore/PostgresIndexStore.js";
+export * from "./kvStore/AzureCosmosMongovCoreKVStore.js";
 export * from "./kvStore/AzureCosmosNoSqlKVStore.js";
 export { PostgresKVStore } from "./kvStore/PostgresKVStore.js";
 export * from "./StorageContext.js";

packages/llamaindex/src/storage/indexStore/AzureCosmosMongovCoreIndexStore.ts

@@ -0,0 +1,49 @@
+import { MongoClient } from "mongodb";
+import { AzureCosmosVCoreKVStore } from "../kvStore/AzureCosmosMongovCoreKVStore.js";
+import { KVIndexStore } from "./KVIndexStore.js";
+
+const DEFAULT_DATABASE = "IndexStoreDB";
+const DEFAULT_COLLECTION = "IndexStoreCollection";
+
+export interface AzureCosmosVCoreIndexStoreArgs {
+  azureCosmosVCoreKVStore: AzureCosmosVCoreKVStore;
+  namespace?: string;
+}
+
+export class AzureCosmosVCoreIndexStore extends KVIndexStore {
+  constructor({
+    azureCosmosVCoreKVStore,
+    namespace,
+  }: AzureCosmosVCoreIndexStoreArgs) {
+    super(azureCosmosVCoreKVStore, namespace);
+  }
+
+  /**
+   * Static method for creating an instance using a MongoClient.
+   * @returns Instance of AzureCosmosVCoreIndexStore
+   * @param mongoClient - MongoClient instance
+   * @param dbName - Database name
+   * @param collectionName - Collection name
+   * @example
+   * ```ts
+   * const mongoClient = new MongoClient("mongodb://localhost:27017");
+   * const indexStore = AzureCosmosVCoreIndexStore.fromMongoClient(mongoClient, "my_db", "my_collection");
+   * ```
+   */
+  static fromMongoClient(
+    mongoClient: MongoClient,
+    dbName: string = DEFAULT_DATABASE,
+    collectionName: string = DEFAULT_COLLECTION,
+  ) {
+    const azureCosmosVCoreKVStore = new AzureCosmosVCoreKVStore({
+      mongoClient,
+      dbName,
+      collectionName,
+    });
+    const namespace = `${dbName}.${collectionName}`;
+    return new AzureCosmosVCoreIndexStore({
+      azureCosmosVCoreKVStore,
+      namespace,
+    });
+  }
+}

packages/llamaindex/src/storage/kvStore/AzureCosmosMongovCoreKVStore.ts

@@ -0,0 +1,88 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import { BaseKVStore } from "@llamaindex/core/storage/kv-store";
+import type { Collection } from "mongodb";
+import { MongoClient } from "mongodb";
+const DEFAULT_CHAT_DATABASE = "KVStoreDB";
+const DEFAULT_CHAT_Collection = "KVStoreCollection";
+
+export interface VcoreConnectionStringOptions
+  extends AzureCosmosVCoreKVStoreConfig {
+  connectionString?: string;
+}
+
+export interface AzureCosmosVCoreKVStoreConfig {
+  mongoClient?: MongoClient;
+  dbName?: string;
+  collectionName?: string;
+}
+
+export class AzureCosmosVCoreKVStore extends BaseKVStore {
+  private mongoClient: MongoClient;
+
+  private dbName: string;
+  private collectionName: string;
+
+  private collection?: Collection;
+
+  /**
+   * Create a new AzureCosmosDBNoSQLVectorStore instance.
+   */
+  constructor({
+    mongoClient,
+    dbName = DEFAULT_CHAT_DATABASE,
+    collectionName = DEFAULT_CHAT_Collection,
+  }: AzureCosmosVCoreKVStoreConfig) {
+    super();
+    if (!mongoClient) {
+      throw new Error(
+        "MongoClient is required for AzureCosmosDBNoSQLVectorStore initialization",
+      );
+    }
+    this.mongoClient = mongoClient;
+    this.dbName = dbName;
+    this.collectionName = collectionName;
+  }
+
+  client(): MongoClient {
+    return this.mongoClient;
+  }
+
+  private async ensureCollection(): Promise<Collection> {
+    if (!this.collection) {
+      this.collection = this.mongoClient
+        .db(this.dbName)
+        .collection(this.collectionName);
+    }
+    return this.collection;
+  }
+
+  async put(key: string, val: Record<string, any>): Promise<void> {
+    const collection = await this.ensureCollection();
+    const insertResult = await collection.insertOne({
+      id: key,
+      messages: val,
+    });
+  }
+
+  async get(key: string): Promise<Record<string, any> | null> {
+    const collection = await this.ensureCollection();
+    const result = await collection.findOne({ id: key });
+    return result || null;
+  }
+
+  async getAll(): Promise<Record<string, Record<string, any>>> {
+    const collection = await this.ensureCollection();
+    const cursor = collection.find();
+    const output: Record<string, Record<string, any>> = {};
+    await cursor.forEach((item) => {
+      output[item.id] = item.messages;
+    });
+    return output;
+  }
+
+  async delete(key: string): Promise<boolean> {
+    const collection = await this.ensureCollection();
+    await collection.deleteOne({ id: key });
+    return true;
+  }
+}