@vectorstores/azure 0.1.0

package/dist/storage.edge-light.js

@@ -0,0 +1,451 @@
+import { BaseKVStore, KVDocumentStore, KVIndexStore } from '@vectorstores/core';
+import { MongoClient } from 'mongodb';
+import { CosmosClient } from '@azure/cosmos';
+import { DefaultAzureCredential } from '@azure/identity';
+import { getEnv } from '@vectorstores/env';
+
+var version = "0.1.0";
+var pkg = {
+	version: version};
+
+const DEFAULT_CHAT_DATABASE$1 = "KVStoreDB";
+const DEFAULT_CHAT_Collection = "KVStoreCollection";
+class AzureCosmosVCoreKVStore extends BaseKVStore {
+    /**
+   * Create a new AzureCosmosDBNoSQLVectorStore instance.
+   */ constructor({ mongoClient, dbName = DEFAULT_CHAT_DATABASE$1, collectionName = DEFAULT_CHAT_Collection }){
+        super();
+        if (!mongoClient) {
+            throw new Error("MongoClient is required for AzureCosmosDBNoSQLVectorStore initialization");
+        }
+        mongoClient.appendMetadata({
+            name: "VECTORSTORES_AZURE_COSMOS_VCORE_KV_STORE",
+            version: pkg.version
+        });
+        this.mongoClient = mongoClient;
+        this.dbName = dbName;
+        this.collectionName = collectionName;
+    }
+    client() {
+        return this.mongoClient;
+    }
+    async ensureCollection() {
+        if (!this.collection) {
+            this.collection = this.mongoClient.db(this.dbName).collection(this.collectionName);
+        }
+        return this.collection;
+    }
+    async put(key, val) {
+        const collection = await this.ensureCollection();
+        await collection.insertOne({
+            id: key,
+            messages: val
+        });
+    }
+    async get(key) {
+        const collection = await this.ensureCollection();
+        const result = await collection.findOne({
+            id: key
+        });
+        return result || null;
+    }
+    async getAll() {
+        const collection = await this.ensureCollection();
+        const cursor = collection.find();
+        const output = {};
+        await cursor.forEach((item)=>{
+            output[item.id] = item.messages;
+        });
+        return output;
+    }
+    async delete(key) {
+        const collection = await this.ensureCollection();
+        await collection.deleteOne({
+            id: key
+        });
+        return true;
+    }
+}
+
+const DEFAULT_DATABASE$3 = "DocumentStoreDB";
+const DEFAULT_COLLECTION$1 = "DocumentStoreCollection";
+class AzureCosmosVCoreDocumentStore extends KVDocumentStore {
+    constructor({ azureCosmosVCoreKVStore, namespace }){
+        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, dbName = DEFAULT_DATABASE$3, collectionName = DEFAULT_COLLECTION$1) {
+        const azureCosmosVCoreKVStore = new AzureCosmosVCoreKVStore({
+            mongoClient,
+            dbName,
+            collectionName
+        });
+        const namespace = `${dbName}.${collectionName}`;
+        return new AzureCosmosVCoreDocumentStore({
+            azureCosmosVCoreKVStore,
+            namespace
+        });
+    }
+    /**
+   * Static method for creating an instance using a connection string.
+   * @returns Instance of AzureCosmosVCoreDocumentStore
+   * @param connectionString - MongoDB connection string
+   * @param dbName - Database name
+   * @param collectionName - Collection name
+   * @example
+   * ```ts
+   * const indexStore = AzureCosmosVCoreDocumentStore.fromConnectionString("mongodb://localhost:27017", "my_db", "my_collection");
+   * ```
+   */ static fromConnectionString(connectionString, dbName = DEFAULT_DATABASE$3, collectionName = DEFAULT_COLLECTION$1) {
+        const mongoClient = new MongoClient(connectionString, {
+            appName: "VECTORSTORES_JS"
+        });
+        return new AzureCosmosVCoreDocumentStore({
+            azureCosmosVCoreKVStore: new AzureCosmosVCoreKVStore({
+                mongoClient,
+                dbName,
+                collectionName
+            }),
+            namespace: `${dbName}.${collectionName}`
+        });
+    }
+}
+
+const USER_AGENT_SUFFIX = "vectorstores-cdbnosql-kvstore-javascript";
+const DEFAULT_CHAT_DATABASE = "KVStoreDB";
+const DEFAULT_CHAT_CONTAINER = "KVStoreContainer";
+const DEFAULT_OFFER_THROUGHPUT = 400;
+function parseConnectionString(connectionString) {
+    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
+    };
+}
+class AzureCosmosNoSqlKVStore extends BaseKVStore {
+    constructor({ cosmosClient, dbName = DEFAULT_CHAT_DATABASE, containerName = DEFAULT_CHAT_CONTAINER, cosmosContainerProperties = {
+        partitionKey: "/id"
+    }, cosmosDatabaseProperties = {} }){
+        super();
+        if (!cosmosClient) {
+            throw new Error("CosmosClient is required for AzureCosmosDBNoSQLKVStore 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() {
+        return this.cosmosClient;
+    }
+    // Asynchronous initialization method to create database and container
+    async init() {
+        // 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 = {}) {
+        const cosmosConnectionString = config.connectionString || getEnv("AZURE_COSMOSDB_NOSQL_CONNECTION_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 = {}) {
+        const cosmosEndpoint = config.endpoint || getEnv("AZURE_COSMOSDB_NOSQL_ACCOUNT_ENDPOINT");
+        const cosmosKey = config.key || getEnv("AZURE_COSMOSDB_NOSQL_ACCOUNT_KEY");
+        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 = {}) {
+        const cosmosEndpoint = config.endpoint || getEnv("AZURE_COSMOSDB_NOSQL_CONNECTION_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, val) {
+        await this.initialize();
+        await this.container.items.upsert({
+            id: key,
+            messages: val
+        });
+    }
+    async get(key) {
+        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() {
+        await this.initialize();
+        const querySpec = {
+            query: "SELECT * from c"
+        };
+        const { resources } = await this.container.items.query(querySpec).fetchAll();
+        const output = resources.reduce((res, item)=>{
+            res[item.id] = item.messages;
+            return res;
+        }, {});
+        return output;
+    }
+    async delete(key) {
+        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;
+        }
+    }
+}
+
+const DEFAULT_DATABASE$2 = "DocumentStoreDB";
+const DEFAULT_CONTAINER$1 = "DocumentStoreContainer";
+class AzureCosmosNoSqlDocumentStore extends KVDocumentStore {
+    constructor({ azureCosmosNoSqlKVStore, namespace }){
+        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 = {}) {
+        options.dbName = options.dbName || DEFAULT_DATABASE$2;
+        options.containerName = options.containerName || DEFAULT_CONTAINER$1;
+        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 = {}) {
+        options.dbName = options.dbName || DEFAULT_DATABASE$2;
+        options.containerName = options.containerName || DEFAULT_CONTAINER$1;
+        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 = {}) {
+        options.dbName = options.dbName || DEFAULT_DATABASE$2;
+        options.containerName = options.containerName || DEFAULT_CONTAINER$1;
+        const azureCosmosNoSqlKVStore = AzureCosmosNoSqlKVStore.fromAadToken(options);
+        const namespace = `${options.dbName}.${options.containerName}`;
+        return new AzureCosmosNoSqlDocumentStore({
+            azureCosmosNoSqlKVStore,
+            namespace
+        });
+    }
+}
+
+const DEFAULT_DATABASE$1 = "IndexStoreDB";
+const DEFAULT_COLLECTION = "IndexStoreCollection";
+class AzureCosmosVCoreIndexStore extends KVIndexStore {
+    constructor({ azureCosmosVCoreKVStore, namespace }){
+        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, dbName = DEFAULT_DATABASE$1, collectionName = DEFAULT_COLLECTION) {
+        const azureCosmosVCoreKVStore = new AzureCosmosVCoreKVStore({
+            mongoClient,
+            dbName,
+            collectionName
+        });
+        const namespace = `${dbName}.${collectionName}`;
+        return new AzureCosmosVCoreIndexStore({
+            azureCosmosVCoreKVStore,
+            namespace
+        });
+    }
+    /**
+   * Static method for creating an instance using a connection string.
+   * @returns Instance of AzureCosmosVCoreIndexStore
+   * @param connectionString - MongoDB connection string
+   * @param dbName - Database name
+   * @param collectionName - Collection name
+   * @example
+   * ```ts
+   * const indexStore = AzureCosmosVCoreIndexStore.fromConnectionString("mongodb://localhost:27017", "my_db", "my_collection");
+   * ```
+   */ static fromConnectionString(connectionString, dbName = DEFAULT_DATABASE$1, collectionName = DEFAULT_COLLECTION) {
+        const mongoClient = new MongoClient(connectionString, {
+            appName: "VECTORSTORES_JS"
+        });
+        return new AzureCosmosVCoreIndexStore({
+            azureCosmosVCoreKVStore: new AzureCosmosVCoreKVStore({
+                mongoClient,
+                dbName,
+                collectionName
+            }),
+            namespace: `${dbName}.${collectionName}`
+        });
+    }
+}
+
+const DEFAULT_DATABASE = "IndexStoreDB";
+const DEFAULT_CONTAINER = "IndexStoreContainer";
+class AzureCosmosNoSqlIndexStore extends KVIndexStore {
+    constructor({ azureCosmosNoSqlKVStore, namespace }){
+        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 = {}) {
+        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 = {}) {
+        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 = {}) {
+        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
+        });
+    }
+}
+
+export { AzureCosmosNoSqlDocumentStore, AzureCosmosNoSqlIndexStore, AzureCosmosNoSqlKVStore, AzureCosmosVCoreDocumentStore, AzureCosmosVCoreIndexStore, AzureCosmosVCoreKVStore };