From efa221116f1676e29e6454c36a27dd23e7b4622c Mon Sep 17 00:00:00 2001
From: fatmelon <708842811@qq.com>
Date: Wed, 20 Nov 2024 09:58:25 +0800
Subject: [PATCH] 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 | 5 ++
.../AzureCosmosMongovCoreDocumentStore.ts | 49 +++++++++++
packages/llamaindex/src/storage/index.ts | 3 +
.../AzureCosmosMongovCoreIndexStore.ts | 49 +++++++++++
.../kvStore/AzureCosmosMongovCoreKVStore.ts | 88 +++++++++++++++++++
5 files changed, 194 insertions(+)
create mode 100644 .changeset/tough-rings-roll.md
create mode 100644 packages/llamaindex/src/storage/docStore/AzureCosmosMongovCoreDocumentStore.ts
create mode 100644 packages/llamaindex/src/storage/indexStore/AzureCosmosMongovCoreIndexStore.ts
create mode 100644 packages/llamaindex/src/storage/kvStore/AzureCosmosMongovCoreKVStore.ts
diff --git a/.changeset/tough-rings-roll.md b/.changeset/tough-rings-roll.md
new file mode 100644
index 000000000..04a577249
--- /dev/null
+++ b/.changeset/tough-rings-roll.md
@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+feat: add Azure Cosmos DB mongo vCore DocumentStore, IndexStore, KVStore
diff --git a/packages/llamaindex/src/storage/docStore/AzureCosmosMongovCoreDocumentStore.ts b/packages/llamaindex/src/storage/docStore/AzureCosmosMongovCoreDocumentStore.ts
new file mode 100644
index 000000000..f9fbe83ff
--- /dev/null
+++ b/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,
+ });
+ }
+}
diff --git a/packages/llamaindex/src/storage/index.ts b/packages/llamaindex/src/storage/index.ts
index bb08b4980..d0b99ae5d 100644
--- a/packages/llamaindex/src/storage/index.ts
+++ b/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";
diff --git a/packages/llamaindex/src/storage/indexStore/AzureCosmosMongovCoreIndexStore.ts b/packages/llamaindex/src/storage/indexStore/AzureCosmosMongovCoreIndexStore.ts
new file mode 100644
index 000000000..2d7118d05
--- /dev/null
+++ b/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,
+ });
+ }
+}
diff --git a/packages/llamaindex/src/storage/kvStore/AzureCosmosMongovCoreKVStore.ts b/packages/llamaindex/src/storage/kvStore/AzureCosmosMongovCoreKVStore.ts
new file mode 100644
index 000000000..534334006
--- /dev/null
+++ b/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;
+ }
+}
--
GitLab