feat: integrate with azure cosmos db (#1444)

Co-authored-by: Wassim Chegham <github@wassim.dev>
Co-authored-by: Alex Yang <himself65@outlook.com>

.changeset/slow-pets-push.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+feat: add Azure Cosmos DB DocumentStore, IndexStore, KVStore, update vectorStore and examples

examples/azure-cosmosdb.ts

@@ -0,0 +1,71 @@
+import "dotenv/config";
+
+import {
+  DefaultAzureCredential,
+  getBearerTokenProvider,
+} from "@azure/identity";
+import {
+  AzureCosmosDBNoSqlVectorStore,
+  AzureCosmosNoSqlDocumentStore,
+  AzureCosmosNoSqlIndexStore,
+  Document,
+  OpenAI,
+  OpenAIEmbedding,
+  Settings,
+  storageContextFromDefaults,
+  VectorStoreIndex,
+} from "llamaindex";
+/**
+ * This example demonstrates how to use Azure CosmosDB with LlamaIndex.
+ * It uses Azure CosmosDB as IndexStore, DocumentStore, and VectorStore.
+ *
+ * To run this example, create an .env file under /examples and set the following environment variables:
+ *
+ * AZURE_OPENAI_ENDPOINT="https://AOAI-ACCOUNT.openai.azure.com" // Sample Azure OpenAI endpoint.
+ * AZURE_DEPLOYMENT_NAME="gpt-4o" // Sample Azure OpenAI deployment name.
+ * EMBEDDING_MODEL="text-embedding-3-large" // Sample Azure OpenAI embedding model.
+ * AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT = "https://DB-ACCOUNT.documents.azure.com:443/" // Sample CosmosDB account endpoint.
+ *
+ * This example uses managed identity to authenticate with Azure CosmosDB and Azure OpenAI. Make sure to assign the required roles to the managed identity.
+ * You can also use connectionString for Azure CosmosDB and Keys with Azure OpenAI for authentication.
+ */
+(async () => {
+  const credential = new DefaultAzureCredential();
+  const azureADTokenProvider = getBearerTokenProvider(
+    credential,
+    "https://cognitiveservices.azure.com/.default",
+  );
+
+  const azure = {
+    azureADTokenProvider,
+    deployment: process.env.AZURE_DEPLOYMENT_NAME,
+  };
+  Settings.llm = new OpenAI({ azure });
+  Settings.embedModel = new OpenAIEmbedding({
+    model: process.env.EMBEDDING_MODEL,
+    azure: {
+      ...azure,
+      deployment: process.env.EMBEDDING_MODEL,
+    },
+  });
+  const docStore = AzureCosmosNoSqlDocumentStore.fromAadToken();
+  console.log({ docStore });
+  const indexStore = AzureCosmosNoSqlIndexStore.fromAadToken();
+  console.log({ indexStore });
+  const vectorStore = AzureCosmosDBNoSqlVectorStore.fromUriAndManagedIdentity();
+  console.log({ vectorStore });
+  const storageContext = await storageContextFromDefaults({
+    docStore,
+    indexStore,
+    vectorStore,
+  });
+  console.log({ storageContext });
+
+  const document = new Document({ text: "Test Text" });
+  const index = await VectorStoreIndex.fromDocuments([document], {
+    storageContext,
+    logProgress: true,
+  });
+
+  console.log({ index });
+})();

examples/cosmosdb/loadVectorData.ts

@@ -6,17 +6,21 @@ import {
 } from "@llamaindex/readers/cosmosdb";
 import * as dotenv from "dotenv";
 import {
-  AzureCosmosDBNoSqlVectorStore,
+  AzureCosmosDBNoSQLConfig,
   OpenAI,
   OpenAIEmbedding,
   Settings,
   storageContextFromDefaults,
   VectorStoreIndex,
 } from "llamaindex";
+import {
+  createStoresFromConnectionString,
+  createStoresFromManagedIdentity,
+} from "./utils";
 // Load environment variables from local .env file
 dotenv.config();
 
-const cosmosEndpoint = process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT!;
+const cosmosEndpoint = process.env.AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT!;
 const cosmosConnectionString =
   process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING!;
 const databaseName =
@@ -26,7 +30,7 @@ const collectionName =
 const vectorCollectionName =
   process.env.AZURE_COSMOSDB_VECTOR_CONTAINER_NAME || "vectorContainer";
 
-// This exampple uses Azure OpenAI llm and embedding models
+// This example uses Azure OpenAI llm and embedding models
 const llmInit = {
   azure: {
     apiVersion: process.env.AZURE_OPENAI_LLM_API_VERSION,
@@ -46,24 +50,48 @@ const embedModelInit = {
 Settings.llm = new OpenAI(llmInit);
 Settings.embedModel = new OpenAIEmbedding(embedModelInit);
 
+// Initialize the CosmosDB client
+async function initializeCosmosClient() {
+  if (cosmosConnectionString) {
+    return new CosmosClient(cosmosConnectionString);
+  } else {
+    const credential = new DefaultAzureCredential();
+    return new CosmosClient({
+      endpoint: cosmosEndpoint,
+      aadCredentials: credential,
+    });
+  }
+}
+
+// Initialize CosmosDB to be used as a vectorStore, docStore, and indexStore
+async function initializeStores() {
+  // Create a configuration object for the Azure CosmosDB NoSQL Vector Store
+  const dbConfig: AzureCosmosDBNoSQLConfig = {
+    databaseName,
+    containerName: vectorCollectionName,
+    flatMetadata: false,
+  };
+
+  if (cosmosConnectionString) {
+    return createStoresFromConnectionString(cosmosConnectionString, dbConfig);
+  } else {
+    // Use managed identity to authenticate with Azure CosmosDB
+    const credential = new DefaultAzureCredential();
+    return createStoresFromManagedIdentity(
+      cosmosEndpoint,
+      credential,
+      dbConfig,
+    );
+  }
+}
+
 async function loadVectorData() {
   if (!cosmosConnectionString && !cosmosEndpoint) {
     throw new Error(
       "Azure CosmosDB connection string or endpoint must be set.",
     );
   }
-
-  let cosmosClient: CosmosClient;
-  // initialize the cosmos client
-  if (cosmosConnectionString) {
-    cosmosClient = new CosmosClient(cosmosConnectionString);
-  } else {
-    cosmosClient = new CosmosClient({
-      endpoint: cosmosEndpoint,
-      aadCredentials: new DefaultAzureCredential(),
-    });
-  }
-
+  const cosmosClient = await initializeCosmosClient();
   const reader = new SimpleCosmosDBReader(cosmosClient);
   // create a configuration object for the reader
   const simpleCosmosReaderConfig: SimpleCosmosDBReaderLoaderConfig = {
@@ -76,16 +104,15 @@ async function loadVectorData() {
 
   // load objects from cosmos and convert them into LlamaIndex Document objects
   const documents = await reader.loadData(simpleCosmosReaderConfig);
-  // create Azure CosmosDB as a vector store
-  const vectorStore = new AzureCosmosDBNoSqlVectorStore({
-    client: cosmosClient,
-    databaseName,
-    containerName: vectorCollectionName,
-    flatMetadata: false,
-  });
 
+  // use Azure CosmosDB as a vectorStore, docStore, and indexStore
+  const { vectorStore, docStore, indexStore } = await initializeStores();
   // Store the embeddings in the CosmosDB container
-  const storageContext = await storageContextFromDefaults({ vectorStore });
+  const storageContext = await storageContextFromDefaults({
+    vectorStore,
+    docStore,
+    indexStore,
+  });
   await VectorStoreIndex.fromDocuments(documents, { storageContext });
   console.log(
     `Successfully created embeddings in the CosmosDB container ${vectorCollectionName}.`,

examples/cosmosdb/queryVectorData.ts

@@ -3,17 +3,21 @@ import { DefaultAzureCredential } from "@azure/identity";
 import * as dotenv from "dotenv";
 import {
   AzureCosmosDBNoSQLConfig,
-  AzureCosmosDBNoSqlVectorStore,
   OpenAI,
   OpenAIEmbedding,
   Settings,
+  storageContextFromDefaults,
   VectorStoreIndex,
 } from "llamaindex";
+import {
+  createStoresFromConnectionString,
+  createStoresFromManagedIdentity,
+} from "./utils";
 
 // Load environment variables from local .env file
 dotenv.config();
 
-const cosmosEndpoint = process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT!;
+const cosmosEndpoint = process.env.AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT!;
 const cosmosConnectionString =
   process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING!;
 const databaseName =
@@ -40,6 +44,27 @@ const embedModelInit = {
 Settings.llm = new OpenAI(llmInit);
 Settings.embedModel = new OpenAIEmbedding(embedModelInit);
 
+async function initializeStores() {
+  // Create a configuration object for the Azure CosmosDB NoSQL Vector Store
+  const dbConfig: AzureCosmosDBNoSQLConfig = {
+    databaseName,
+    containerName,
+    flatMetadata: false,
+  };
+
+  if (cosmosConnectionString) {
+    return createStoresFromConnectionString(cosmosConnectionString, dbConfig);
+  } else {
+    // Use managed identity to authenticate with Azure CosmosDB
+    const credential = new DefaultAzureCredential();
+    return createStoresFromManagedIdentity(
+      cosmosEndpoint,
+      credential,
+      dbConfig,
+    );
+  }
+}
+
 async function query() {
   if (!cosmosConnectionString && !cosmosEndpoint) {
     throw new Error(
@@ -65,10 +90,19 @@ async function query() {
     containerName,
     flatMetadata: false,
   };
-  const store = new AzureCosmosDBNoSqlVectorStore(dbConfig);
+
+  // use Azure CosmosDB as a vectorStore, docStore, and indexStore
+  const { vectorStore, docStore, indexStore } = await initializeStores();
+
+  // Store the embeddings in the CosmosDB container
+  const storageContext = await storageContextFromDefaults({
+    vectorStore,
+    docStore,
+    indexStore,
+  });
 
   // create an index from the Azure CosmosDB NoSQL Vector Store
-  const index = await VectorStoreIndex.fromVectorStore(store);
+  const index = await VectorStoreIndex.init({ storageContext });
 
   // create a retriever and a query engine from the index
   const retriever = index.asRetriever({ similarityTopK: 20 });

examples/cosmosdb/utils.ts

@@ -0,0 +1,51 @@
+import { TokenCredential } from "@azure/identity";
+import {
+  AzureCosmosDBNoSQLConfig,
+  AzureCosmosDBNoSqlVectorStore,
+  AzureCosmosNoSqlDocumentStore,
+  AzureCosmosNoSqlIndexStore,
+} from "llamaindex";
+
+/**
+ * Util function to create AzureCosmosDB vectorStore, docStore, indexStore from connection string.
+ */
+export const createStoresFromConnectionString = (
+  connectionString: string,
+  dbConfig: AzureCosmosDBNoSQLConfig,
+) => {
+  const vectorStore = AzureCosmosDBNoSqlVectorStore.fromConnectionString({
+    connectionString,
+    ...dbConfig,
+  });
+  const docStore = AzureCosmosNoSqlDocumentStore.fromConnectionString({
+    connectionString,
+  });
+  const indexStore = AzureCosmosNoSqlIndexStore.fromConnectionString({
+    connectionString,
+  });
+  return { vectorStore, docStore, indexStore };
+};
+
+/**
+ * Util function to create AzureCosmosDB vectorStore, docStore, indexStore from connection string.
+ */
+export const createStoresFromManagedIdentity = (
+  endpoint: string,
+  credential: TokenCredential,
+  dbConfig: AzureCosmosDBNoSQLConfig,
+) => {
+  const vectorStore = AzureCosmosDBNoSqlVectorStore.fromUriAndManagedIdentity({
+    endpoint,
+    credential,
+    ...dbConfig,
+  });
+  const docStore = AzureCosmosNoSqlDocumentStore.fromAadToken({
+    endpoint,
+    credential,
+  });
+  const indexStore = AzureCosmosNoSqlIndexStore.fromAadToken({
+    endpoint,
+    credential,
+  });
+  return { vectorStore, docStore, indexStore };
+};

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

@@ -0,0 +1,77 @@
+import {
+  AzureCosmosNoSqlKVStore,
+  type AadTokenOptions,
+  type AccountAndKeyOptions,
+  type ConnectionStringOptions,
+} from "../kvStore/AzureCosmosNoSqlKVStore.js";
+import { KVDocumentStore } from "./KVDocumentStore.js";
+
+const DEFAULT_DATABASE = "DocumentStoreDB";
+const DEFAULT_CONTAINER = "DocumentStoreContainer";
+
+export interface AzureCosmosNoSqlDocumentStoreArgs {
+  azureCosmosNoSqlKVStore: AzureCosmosNoSqlKVStore;
+  namespace?: string;
+}
+
+export class AzureCosmosNoSqlDocumentStore extends KVDocumentStore {
+  constructor({
+    azureCosmosNoSqlKVStore,
+    namespace,
+  }: AzureCosmosNoSqlDocumentStoreArgs) {
+    super(azureCosmosNoSqlKVStore, namespace);
+  }
+
+  /**
+   * Static method for creating an instance using a connection string.
+   * If no connection string is provided, it will attempt to use the env variable `AZURE_COSMOSDB_NOSQL_CONNECTION_STRING` as connection string.
+   * @returns Instance of AzureCosmosNoSqlDocumentStore
+   */
+  static fromConnectionString(options: ConnectionStringOptions = {}) {
+    options.dbName = options.dbName || DEFAULT_DATABASE;
+    options.containerName = options.containerName || DEFAULT_CONTAINER;
+
+    const azureCosmosNoSqlKVStore =
+      AzureCosmosNoSqlKVStore.fromConnectionString(options);
+    const namespace = `${options.dbName}.${options.containerName}`;
+    return new AzureCosmosNoSqlDocumentStore({
+      azureCosmosNoSqlKVStore,
+      namespace,
+    });
+  }
+
+  /**
+   * Static method for creating an instance using a account endpoint and key.
+   * If no endpoint and key  is provided, it will attempt to use the env variable `AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT` as enpoint and `AZURE_COSMOSDB_NOSQL_ACCOUNT_KEY` as key.
+   * @returns Instance of AzureCosmosNoSqlDocumentStore
+   */
+  static fromAccountAndKey(options: AccountAndKeyOptions = {}) {
+    options.dbName = options.dbName || DEFAULT_DATABASE;
+    options.containerName = options.containerName || DEFAULT_CONTAINER;
+
+    const azureCosmosNoSqlKVStore =
+      AzureCosmosNoSqlKVStore.fromAccountAndKey(options);
+    const namespace = `${options.dbName}.${options.containerName}`;
+    return new AzureCosmosNoSqlDocumentStore({
+      azureCosmosNoSqlKVStore,
+      namespace,
+    });
+  }
+
+  /**
+   * Static method for creating an instance using AAD token.
+   * If no endpoint and credentials are provided, it will attempt to use the env variable `AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT` as endpoint and use DefaultAzureCredential() as credentials.
+   * @returns Instance of AzureCosmosNoSqlDocumentStore
+   */
+  static fromAadToken(options: AadTokenOptions = {}) {
+    options.dbName = options.dbName || DEFAULT_DATABASE;
+    options.containerName = options.containerName || DEFAULT_CONTAINER;
+    const azureCosmosNoSqlKVStore =
+      AzureCosmosNoSqlKVStore.fromAadToken(options);
+    const namespace = `${options.dbName}.${options.containerName}`;
+    return new AzureCosmosNoSqlDocumentStore({
+      azureCosmosNoSqlKVStore,
+      namespace,
+    });
+  }
+}

packages/llamaindex/src/storage/index.ts

@@ -1,11 +1,14 @@
 export * from "@llamaindex/core/storage/chat-store";
+export * from "./docStore/AzureCosmosNoSqlDocumentStore.js";
 export { PostgresDocumentStore } from "./docStore/PostgresDocumentStore.js";
 export { SimpleDocumentStore } from "./docStore/SimpleDocumentStore.js";
 export * from "./docStore/types.js";
 export * from "./FileSystem.js";
+export * from "./indexStore/AzureCosmosNoSqlIndexStore.js";
 export { PostgresIndexStore } from "./indexStore/PostgresIndexStore.js";
 export { SimpleIndexStore } from "./indexStore/SimpleIndexStore.js";
 export * from "./indexStore/types.js";
+export * from "./kvStore/AzureCosmosNoSqlKVStore.js";
 export { PostgresKVStore } from "./kvStore/PostgresKVStore.js";
 export { SimpleKVStore } from "./kvStore/SimpleKVStore.js";
 export * from "./kvStore/types.js";

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

@@ -0,0 +1,77 @@
+import {
+  AzureCosmosNoSqlKVStore,
+  type AadTokenOptions,
+  type AccountAndKeyOptions,
+  type ConnectionStringOptions,
+} from "../kvStore/AzureCosmosNoSqlKVStore.js";
+import { KVIndexStore } from "./KVIndexStore.js";
+
+const DEFAULT_DATABASE = "IndexStoreDB";
+const DEFAULT_CONTAINER = "IndexStoreContainer";
+
+export interface AzureCosmosNoSqlIndexStoreArgs {
+  azureCosmosNoSqlKVStore: AzureCosmosNoSqlKVStore;
+  namespace?: string;
+}
+
+export class AzureCosmosNoSqlIndexStore extends KVIndexStore {
+  constructor({
+    azureCosmosNoSqlKVStore,
+    namespace,
+  }: AzureCosmosNoSqlIndexStoreArgs) {
+    super(azureCosmosNoSqlKVStore, namespace);
+  }
+
+  /**
+   * Static method for creating an instance using a connection string.
+   * If no connection string is provided, it will attempt to use the env variable `AZURE_COSMOSDB_NOSQL_CONNECTION_STRING` as connection string.
+   * @returns Instance of AzureCosmosNoSqlIndexStore
+   */
+  static fromConnectionString(options: ConnectionStringOptions = {}) {
+    options.dbName = options.dbName || DEFAULT_DATABASE;
+    options.containerName = options.containerName || DEFAULT_CONTAINER;
+    const azureCosmosNoSqlKVStore =
+      AzureCosmosNoSqlKVStore.fromConnectionString(options);
+    const namespace = `${options.dbName}.${options.containerName}`;
+    return new AzureCosmosNoSqlIndexStore({
+      azureCosmosNoSqlKVStore,
+      namespace,
+    });
+  }
+
+  /**
+   * Static method for creating an instance using a account endpoint and key.
+   * If no endpoint and key  is provided, it will attempt to use the env variable `AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT` as enpoint and `AZURE_COSMOSDB_NOSQL_ACCOUNT_KEY` as key.
+   * @returns Instance of AzureCosmosNoSqlIndexStore
+   */
+  static fromAccountAndKey(options: AccountAndKeyOptions = {}) {
+    options.dbName = options.dbName || DEFAULT_DATABASE;
+    options.containerName = options.containerName || DEFAULT_CONTAINER;
+
+    const azureCosmosNoSqlKVStore =
+      AzureCosmosNoSqlKVStore.fromAccountAndKey(options);
+    const namespace = `${options.dbName}.${options.containerName}`;
+    return new AzureCosmosNoSqlIndexStore({
+      azureCosmosNoSqlKVStore,
+      namespace,
+    });
+  }
+
+  /**
+   * Static method for creating an instance using AAD token.
+   * If no endpoint and credentials are provided, it will attempt to use the env variable `AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT` as endpoint and use DefaultAzureCredential() as credentials.
+   * @returns Instance of AzureCosmosNoSqlIndexStore
+   */
+  static fromAadToken(options: AadTokenOptions = {}) {
+    options.dbName = options.dbName || DEFAULT_DATABASE;
+    options.containerName = options.containerName || DEFAULT_CONTAINER;
+
+    const azureCosmosNoSqlKVStore =
+      AzureCosmosNoSqlKVStore.fromAadToken(options);
+    const namespace = `${options.dbName}.${options.containerName}`;
+    return new AzureCosmosNoSqlIndexStore({
+      azureCosmosNoSqlKVStore,
+      namespace,
+    });
+  }
+}

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

@@ -0,0 +1,274 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import { Container, CosmosClient, Database } from "@azure/cosmos";
+import { DefaultAzureCredential, type TokenCredential } from "@azure/identity";
+import { getEnv } from "@llamaindex/env";
+import { BaseKVStore } from "./types.js";
+const USER_AGENT_SUFFIX = "LlamaIndex-CDBNoSQL-KVStore-JavaScript";
+const DEFAULT_CHAT_DATABASE = "KVStoreDB";
+const DEFAULT_CHAT_CONTAINER = "KVStoreContainer";
+const DEFAULT_OFFER_THROUGHPUT = 400;
+
+function parseConnectionString(connectionString: string): {
+  endpoint: string;
+  key: string;
+} {
+  const parts = connectionString.split(";");
+  let endpoint = "";
+  let accountKey = "";
+
+  parts.forEach((part) => {
+    const [key, value] = part.split("=");
+    if (key && key.trim() === "AccountEndpoint") {
+      endpoint = value?.trim() ?? "";
+    } else if ((key ?? "").trim() === "AccountKey") {
+      accountKey = value?.trim() ?? "";
+    }
+  });
+
+  if (!endpoint || !accountKey) {
+    throw new Error(
+      "Invalid connection string: missing AccountEndpoint or AccountKey.",
+    );
+  }
+  return { endpoint, key: accountKey };
+}
+
+export interface CosmosDatabaseProperties {
+  throughput?: number;
+}
+
+export interface CosmosContainerProperties {
+  partitionKey: any;
+  [key: string]: any;
+}
+export interface ConnectionStringOptions extends AzureCosmosNoSqlKVStoreConfig {
+  connectionString?: string;
+}
+export interface AccountAndKeyOptions extends AzureCosmosNoSqlKVStoreConfig {
+  endpoint?: string;
+  key?: string;
+}
+export interface AadTokenOptions extends AzureCosmosNoSqlKVStoreConfig {
+  endpoint?: string;
+  credential?: TokenCredential;
+}
+export interface AzureCosmosNoSqlKVStoreConfig {
+  cosmosClient?: CosmosClient;
+  dbName?: string;
+  containerName?: string;
+  cosmosContainerProperties?: CosmosContainerProperties;
+  cosmosDatabaseProperties?: CosmosDatabaseProperties;
+}
+
+export class AzureCosmosNoSqlKVStore extends BaseKVStore {
+  private cosmosClient: CosmosClient;
+  private database!: Database;
+  private container!: Container;
+  private initPromise?: Promise<void>;
+
+  private dbName: string;
+  private containerName: string;
+  private cosmosContainerProperties: CosmosContainerProperties;
+  private cosmosDatabaseProperties: CosmosDatabaseProperties;
+  private initialize: () => Promise<void>;
+
+  constructor({
+    cosmosClient,
+    dbName = DEFAULT_CHAT_DATABASE,
+    containerName = DEFAULT_CHAT_CONTAINER,
+    cosmosContainerProperties = { partitionKey: "/id" },
+    cosmosDatabaseProperties = {},
+  }: AzureCosmosNoSqlKVStoreConfig) {
+    super();
+    if (!cosmosClient) {
+      throw new Error(
+        "CosmosClient is required for AzureCosmosDBNoSQLVectorStore initialization",
+      );
+    }
+    this.cosmosClient = cosmosClient;
+    this.dbName = dbName;
+    this.containerName = containerName;
+    this.cosmosContainerProperties = cosmosContainerProperties;
+    this.cosmosDatabaseProperties = cosmosDatabaseProperties;
+
+    this.initialize = () => {
+      if (this.initPromise === undefined) {
+        this.initPromise = this.init().catch((error) => {
+          console.error(
+            "Error during AzureCosmosDBNoSQLKVStore initialization",
+            error,
+          );
+        });
+      }
+      return this.initPromise;
+    };
+  }
+
+  client(): CosmosClient {
+    return this.cosmosClient;
+  }
+
+  // Asynchronous initialization method to create database and container
+  private async init(): Promise<void> {
+    // Set default throughput if not provided
+    const throughput =
+      this.cosmosDatabaseProperties?.throughput || DEFAULT_OFFER_THROUGHPUT;
+
+    // Create the database if it doesn't exist
+    const { database } = await this.cosmosClient.databases.createIfNotExists({
+      id: this.dbName,
+      throughput,
+    });
+    this.database = database;
+
+    // Create the container if it doesn't exist
+    const { container } = await this.database.containers.createIfNotExists({
+      id: this.containerName,
+      throughput: this.cosmosContainerProperties?.throughput,
+      partitionKey: this.cosmosContainerProperties?.partitionKey,
+      indexingPolicy: this.cosmosContainerProperties?.indexingPolicy,
+      defaultTtl: this.cosmosContainerProperties?.defaultTtl,
+      uniqueKeyPolicy: this.cosmosContainerProperties?.uniqueKeyPolicy,
+      conflictResolutionPolicy:
+        this.cosmosContainerProperties?.conflictResolutionPolicy,
+      computedProperties: this.cosmosContainerProperties?.computedProperties,
+    });
+    this.container = container;
+  }
+  /**
+   * Static method for creating an instance using a connection string.
+   * If no connection string is provided, it will attempt to use the env variable `AZURE_COSMOSDB_NOSQL_CONNECTION_STRING` as connection string.
+   * @returns Instance of AzureCosmosNoSqlKVStore
+   */
+  static fromConnectionString(
+    config: { connectionString?: string } & AzureCosmosNoSqlKVStoreConfig = {},
+  ): AzureCosmosNoSqlKVStore {
+    const cosmosConnectionString =
+      config.connectionString ||
+      (getEnv("AZURE_COSMOSDB_NOSQL_CONNECTION_STRING") as string);
+    if (!cosmosConnectionString) {
+      throw new Error("Azure CosmosDB connection string must be provided");
+    }
+    const { endpoint, key } = parseConnectionString(cosmosConnectionString);
+    const cosmosClient = new CosmosClient({
+      endpoint,
+      key,
+      userAgentSuffix: USER_AGENT_SUFFIX,
+    });
+    return new AzureCosmosNoSqlKVStore({
+      ...config,
+      cosmosClient,
+    });
+  }
+
+  /**
+   * Static method for creating an instance using a account endpoint and key.
+   * If no endpoint and key  is provided, it will attempt to use the env variable `AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT` as enpoint and `AZURE_COSMOSDB_NOSQL_ACCOUNT_KEY` as key.
+   * @returns Instance of AzureCosmosNoSqlKVStore
+   */
+  static fromAccountAndKey(
+    config: {
+      endpoint?: string;
+      key?: string;
+    } & AzureCosmosNoSqlKVStoreConfig = {},
+  ): AzureCosmosNoSqlKVStore {
+    const cosmosEndpoint =
+      config.endpoint ||
+      (getEnv("AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT") as string);
+    const cosmosKey =
+      config.key || (getEnv("AZURE_COSMOSDB_NOSQL_ACCOUNT_KEY") as string);
+
+    if (!cosmosEndpoint || !cosmosKey) {
+      throw new Error(
+        "Azure CosmosDB account endpoint and key must be provided",
+      );
+    }
+    const cosmosClient = new CosmosClient({
+      endpoint: cosmosEndpoint,
+      key: cosmosKey,
+      userAgentSuffix: USER_AGENT_SUFFIX,
+    });
+    return new AzureCosmosNoSqlKVStore({
+      ...config,
+      cosmosClient,
+    });
+  }
+
+  /**
+   * Static method for creating an instance using AAD token.
+   * If no endpoint and credentials are provided, it will attempt to use the env variable `AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT` as endpoint and use DefaultAzureCredential() as credentials.
+   * @returns Instance of AzureCosmosNoSqlKVStore
+   */
+  static fromAadToken(
+    config: {
+      endpoint?: string;
+      credential?: TokenCredential;
+    } & AzureCosmosNoSqlKVStoreConfig = {},
+  ): AzureCosmosNoSqlKVStore {
+    const cosmosEndpoint =
+      config.endpoint ||
+      (getEnv("AZURE_COSMOSDB_NOSQL_CONNECTION_STRING") as string);
+
+    if (!cosmosEndpoint) {
+      throw new Error("Azure CosmosDB account endpoint must be provided");
+    }
+    const credentials = config.credential ?? new DefaultAzureCredential();
+    const cosmosClient = new CosmosClient({
+      endpoint: cosmosEndpoint,
+      aadCredentials: credentials,
+      userAgentSuffix: USER_AGENT_SUFFIX,
+    });
+    return new AzureCosmosNoSqlKVStore({
+      ...config,
+      cosmosClient,
+    });
+  }
+
+  async put(key: string, val: Record<string, any>): Promise<void> {
+    await this.initialize();
+    await this.container.items.upsert({
+      id: key,
+      messages: val,
+    });
+  }
+
+  async get(key: string): Promise<Record<string, any> | null> {
+    await this.initialize();
+    try {
+      const { resource } = await this.container.item(key).read();
+      return resource?.messages || null;
+    } catch (error) {
+      console.error(`Error retrieving item with key ${key}:`, error);
+      return null;
+    }
+  }
+
+  async getAll(): Promise<Record<string, Record<string, any>>> {
+    await this.initialize();
+    const querySpec = {
+      query: "SELECT * from c",
+    };
+    const { resources } = await this.container.items
+      .query(querySpec)
+      .fetchAll();
+    const output: Record<string, Record<string, any>> = resources.reduce(
+      (res, item) => {
+        res[item.id] = item.messages;
+        return res;
+      },
+      {},
+    );
+    return output;
+  }
+
+  async delete(key: string): Promise<boolean> {
+    await this.initialize();
+    try {
+      await this.container.item(key).delete();
+      return true;
+    } catch (error) {
+      console.error(`Error deleting item with key ${key}:`, error);
+      return false;
+    }
+  }
+}

packages/llamaindex/src/vector-store/AzureCosmosDBNoSqlVectorStore.ts

@@ -5,7 +5,6 @@ import {
   VectorEmbeddingDistanceFunction,
   VectorIndexType,
   type ContainerRequest,
-  type CosmosClientOptions,
   type DatabaseRequest,
   type IndexingPolicy,
   type VectorEmbeddingPolicy,
@@ -48,10 +47,7 @@ export interface AzureCosmosDBNoSQLInitOptions {
  */
 export interface AzureCosmosDBNoSQLConfig
   extends AzureCosmosDBNoSQLInitOptions {
-  readonly client?: CosmosClient;
-  readonly connectionString?: string;
-  readonly endpoint?: string;
-  readonly credentials?: TokenCredential;
+  client?: CosmosClient;
   readonly databaseName?: string;
   readonly containerName?: string;
   readonly textKey?: string;
@@ -60,7 +56,7 @@ export interface AzureCosmosDBNoSQLConfig
   readonly idKey?: string;
 }
 
-const USER_AGENT_PREFIX = "LlamaIndex-CDBNoSQL-VectorStore-JavaScript";
+const USER_AGENT_SUFFIX = "LlamaIndex-CDBNoSQL-VectorStore-JavaScript";
 
 const DEFAULT_VECTOR_EMBEDDING_POLICY = {
   vectorEmbeddings: [
@@ -148,40 +144,12 @@ export class AzureCosmosDBNoSqlVectorStore extends BaseVectorStore {
 
   constructor(dbConfig: AzureCosmosDBNoSQLConfig & VectorStoreBaseParams) {
     super(dbConfig);
-    const connectionString =
-      dbConfig.connectionString ??
-      getEnv("AZURE_COSMOSDB_NOSQL_CONNECTION_STRING");
-
-    const endpoint =
-      dbConfig.endpoint ?? getEnv("AZURE_COSMOSDB_NOSQL_ENDPOINT");
-
-    if (!dbConfig.client && !connectionString && !endpoint) {
+    if (!dbConfig.client) {
       throw new Error(
-        "CosmosDB client, connection string or endpoint must be set in the configuration.",
+        "CosmosClient is required for AzureCosmosDBNoSQLVectorStore initialization",
       );
     }
-
-    if (!dbConfig.client) {
-      if (connectionString) {
-        const { endpoint, key } = parseConnectionString(connectionString);
-        this.cosmosClient = new CosmosClient({
-          endpoint,
-          key,
-          userAgentSuffix: USER_AGENT_PREFIX,
-        } as CosmosClientOptions);
-      } else {
-        // Use managed identity
-        this.cosmosClient = new CosmosClient({
-          endpoint,
-          aadCredentials: dbConfig.credentials ?? new DefaultAzureCredential(),
-          userAgentSuffix: USER_AGENT_PREFIX,
-        } as CosmosClientOptions);
-      }
-    } else {
-      this.cosmosClient = dbConfig.client;
-    }
-
-    const client = this.cosmosClient;
+    this.cosmosClient = dbConfig.client;
     const databaseName = dbConfig.databaseName ?? "vectorSearchDB";
     const containerName = dbConfig.containerName ?? "vectorSearchContainer";
     this.idKey = dbConfig.idKey ?? "id";
@@ -206,12 +174,17 @@ export class AzureCosmosDBNoSqlVectorStore extends BaseVectorStore {
     // Deferring initialization to the first call to `initialize`
     this.initialize = () => {
       if (this.initPromise === undefined) {
-        this.initPromise = this.init(client, databaseName, containerName, {
-          vectorEmbeddingPolicy,
-          indexingPolicy,
-          createContainerOptions: dbConfig.createContainerOptions,
-          createDatabaseOptions: dbConfig.createDatabaseOptions,
-        }).catch((error) => {
+        this.initPromise = this.init(
+          this.cosmosClient,
+          databaseName,
+          containerName,
+          {
+            vectorEmbeddingPolicy,
+            indexingPolicy,
+            createContainerOptions: dbConfig.createContainerOptions,
+            createDatabaseOptions: dbConfig.createDatabaseOptions,
+          },
+        ).catch((error) => {
           console.error(
             "Error during AzureCosmosDBNoSQLVectorStore initialization",
             error,
@@ -222,6 +195,84 @@ export class AzureCosmosDBNoSqlVectorStore extends BaseVectorStore {
     };
   }
 
+  /**
+   * Static method for creating an instance using a connection string.
+   * If no connection string is provided, it will attempt to use the env variable `AZURE_COSMOSDB_NOSQL_CONNECTION_STRING` as connection string.
+   * @returns Instance of AzureCosmosDBNoSqlVectorStore
+   */
+  static fromConnectionString(
+    config: { connectionString?: string } & AzureCosmosDBNoSQLConfig &
+      VectorStoreBaseParams = {},
+  ): AzureCosmosDBNoSqlVectorStore {
+    const cosmosConnectionString =
+      config.connectionString ||
+      (getEnv("AZURE_COSMOSDB_NOSQL_CONNECTION_STRING") as string);
+    if (!cosmosConnectionString) {
+      throw new Error("Azure CosmosDB connection string must be provided");
+    }
+    const { endpoint, key } = parseConnectionString(cosmosConnectionString);
+    const client = new CosmosClient({
+      endpoint,
+      key,
+      userAgentSuffix: USER_AGENT_SUFFIX,
+    });
+    return new AzureCosmosDBNoSqlVectorStore({ ...config, client });
+  }
+
+  /**
+   * Static method for creating an instance using a account endpoint and key.
+   * If no endpoint and key  is provided, it will attempt to use the env variable `AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT` as enpoint and `AZURE_COSMOSDB_NOSQL_ACCOUNT_KEY` as key.
+   * @returns Instance of AzureCosmosDBNoSqlVectorStore
+   */
+  static fromAccountAndKey(
+    config: { endpoint?: string; key?: string } & AzureCosmosDBNoSQLConfig &
+      VectorStoreBaseParams = {},
+  ): AzureCosmosDBNoSqlVectorStore {
+    const cosmosEndpoint =
+      config.endpoint || (getEnv("AZURE_COSMOSDB_NOSQL_ENDPOINT") as string);
+
+    const cosmosKey =
+      config.key || (getEnv("AZURE_COSMOSDB_NOSQL_KEY") as string);
+
+    if (!cosmosEndpoint || !cosmosKey) {
+      throw new Error(
+        "Azure CosmosDB account endpoint and key must be provided",
+      );
+    }
+    const client = new CosmosClient({
+      endpoint: cosmosEndpoint,
+      key: cosmosKey,
+      userAgentSuffix: USER_AGENT_SUFFIX,
+    });
+    return new AzureCosmosDBNoSqlVectorStore({ ...config, client });
+  }
+
+  /**
+   * Static method for creating an instance using account endpoint and managed identity.
+   * If no endpoint and credentials are provided, it will attempt to use the env variable `AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT` as endpoint and use DefaultAzureCredential() as credentials.
+   * @returns Instance of AzureCosmosDBNoSqlVectorStore
+   */
+  static fromUriAndManagedIdentity(
+    config: {
+      endpoint?: string;
+      credential?: TokenCredential;
+    } & AzureCosmosDBNoSQLConfig &
+      VectorStoreBaseParams = {},
+  ): AzureCosmosDBNoSqlVectorStore {
+    const cosmosEndpoint =
+      config.endpoint ||
+      (getEnv("AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT") as string);
+    if (!cosmosEndpoint) {
+      throw new Error("Azure CosmosDB account endpoint must be provided");
+    }
+    const credentials = config.credential ?? new DefaultAzureCredential();
+    const client = new CosmosClient({
+      endpoint: cosmosEndpoint,
+      aadCredentials: credentials,
+      userAgentSuffix: USER_AGENT_SUFFIX,
+    });
+    return new AzureCosmosDBNoSqlVectorStore({ ...config, client });
+  }
   /**
    * Adds document to the CosmosDB container.
    *