diff --git a/libs/langchain-azure-cosmosdb/src/chat_histories/mongodb.ts b/libs/langchain-azure-cosmosdb/src/chat_histories/mongodb.ts new file mode 100644 index 000000000000..53104c198d71 --- /dev/null +++ b/libs/langchain-azure-cosmosdb/src/chat_histories/mongodb.ts @@ -0,0 +1,155 @@ +import { + Collection, + Document as AzureCosmosMongoDBDocument, + PushOperator, + Db, + MongoClient, +} from "mongodb"; +import { BaseListChatMessageHistory } from "@langchain/core/chat_history"; +import { + BaseMessage, + mapChatMessagesToStoredMessages, + mapStoredMessagesToChatMessages, +} from "@langchain/core/messages"; +import { getEnvironmentVariable } from "@langchain/core/utils/env"; + +export interface AzureCosmosDBMongoChatHistoryDBConfig { + readonly client?: MongoClient; + readonly connectionString?: string; + readonly databaseName?: string; + readonly collectionName?: string; +} + +const ID_KEY = "sessionId"; + +export class AzureCosmosDBMongoChatMessageHistory extends BaseListChatMessageHistory { + lc_namespace = ["langchain", "stores", "message", "azurecosmosdb"]; + + get lc_secrets(): { [key: string]: string } { + return { + connectionString: "AZURE_COSMOSDB_MONGODB_CONNECTION_STRING", + }; + } + + private initPromise?: Promise; + + private readonly client: MongoClient | undefined; + + private database: Db; + + private collection: Collection; + + private sessionId: string; + + initialize: () => Promise; + + constructor( + dbConfig: AzureCosmosDBMongoChatHistoryDBConfig, + sessionId: string + ) { + super(); + + const connectionString = + dbConfig.connectionString ?? + getEnvironmentVariable("AZURE_COSMOSDB_MONGODB_CONNECTION_STRING"); + + if (!dbConfig.client && !connectionString) { + throw new Error("Mongo client or connection string must be set."); + } + + if (!dbConfig.client) { + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + this.client = new MongoClient(connectionString!, { + appName: "langchainjs", + }); + } + + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + const client = dbConfig.client || this.client!; + const databaseName = dbConfig.databaseName ?? "chatHistoryDB"; + const collectionName = dbConfig.collectionName ?? "chatHistory"; + + this.sessionId = sessionId; + + // Deferring initialization to the first call to `initialize` + this.initialize = () => { + if (this.initPromise === undefined) { + this.initPromise = this.init( + client, + databaseName, + collectionName + ).catch((error) => { + console.error( + "Error during AzureCosmosDBMongoChatMessageHistory initialization: ", + error + ); + }); + } + + return this.initPromise; + }; + } + + /** + * Initializes the AzureCosmosDBMongoChatMessageHistory by connecting to the database. + * @param client The MongoClient to use for connecting to the database. + * @param databaseName The name of the database to use. + * @param collectionName The name of the collection to use. + * @returns A promise that resolves when the AzureCosmosDBMongoChatMessageHistory has been initialized. + */ + private async init( + client: MongoClient, + databaseName: string, + collectionName: string + ): Promise { + this.initPromise = (async () => { + await client.connect(); + this.database = client.db(databaseName); + this.collection = this.database.collection(collectionName); + })(); + + return this.initPromise; + } + + /** + * Retrieves the messages stored in the history. + * @returns A promise that resolves with the messages stored in the history. + */ + async getMessages(): Promise { + await this.initialize(); + + const document = await this.collection.findOne({ + [ID_KEY]: this.sessionId, + }); + const messages = document?.messages || []; + return mapStoredMessagesToChatMessages(messages); + } + + /** + * Adds a message to the history. + * @param message The message to add to the history. + * @returns A promise that resolves when the message has been added to the history. + */ + async addMessage(message: BaseMessage): Promise { + await this.initialize(); + + const messages = mapChatMessagesToStoredMessages([message]); + await this.collection.updateOne( + { [ID_KEY]: this.sessionId }, + { + $push: { messages: { $each: messages } } as PushOperator, + }, + { upsert: true } + ); + } + + /** + * Clear the history. + * @returns A promise that resolves when the history has been cleared. + */ + async clear(): Promise { + await this.initialize(); + + await this.collection.deleteOne({ [ID_KEY]: this.sessionId }); + } +} diff --git a/libs/langchain-azure-cosmosdb/src/chat_histories.ts b/libs/langchain-azure-cosmosdb/src/chat_histories/nosql.ts similarity index 100% rename from libs/langchain-azure-cosmosdb/src/chat_histories.ts rename to libs/langchain-azure-cosmosdb/src/chat_histories/nosql.ts diff --git a/libs/langchain-azure-cosmosdb/src/index.ts b/libs/langchain-azure-cosmosdb/src/index.ts index c5160397b474..883710c842ff 100644 --- a/libs/langchain-azure-cosmosdb/src/index.ts +++ b/libs/langchain-azure-cosmosdb/src/index.ts @@ -1,4 +1,5 @@ export * from "./azure_cosmosdb_mongodb.js"; export * from "./azure_cosmosdb_nosql.js"; export * from "./caches.js"; -export * from "./chat_histories.js"; +export * from "./chat_histories/nosql.js"; +export * from "./chat_histories/mongodb.js"; diff --git a/libs/langchain-azure-cosmosdb/src/tests/chat_histories/mongodb.int.test.ts b/libs/langchain-azure-cosmosdb/src/tests/chat_histories/mongodb.int.test.ts new file mode 100644 index 000000000000..35c4a2cf0311 --- /dev/null +++ b/libs/langchain-azure-cosmosdb/src/tests/chat_histories/mongodb.int.test.ts @@ -0,0 +1,95 @@ +/* eslint-disable no-process-env */ + +import { MongoClient, ObjectId } from "mongodb"; +import { AIMessage, HumanMessage } from "@langchain/core/messages"; +import { + AzureCosmosDBMongoChatMessageHistory, + AzureCosmosDBMongoChatHistoryDBConfig, +} from "../../chat_histories/mongodb.js"; + +afterAll(async () => { + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + const client = new MongoClient( + process.env.AZURE_COSMOSDB_MONGODB_CONNECTION_STRING! + ); + await client.connect(); + await client.db("langchain").dropDatabase(); + await client.close(); +}); + +test("Test Azure Cosmos MongoDB history store", async () => { + expect(process.env.AZURE_COSMOSDB_MONGODB_CONNECTION_STRING).toBeDefined(); + + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + const mongoClient = new MongoClient( + process.env.AZURE_COSMOSDB_MONGODB_CONNECTION_STRING! + ); + const dbcfg: AzureCosmosDBMongoChatHistoryDBConfig = { + client: mongoClient, + connectionString: process.env.AZURE_COSMOSDB_MONGODB_CONNECTION_STRING, + databaseName: "langchain", + collectionName: "chathistory", + }; + + const sessionId = new ObjectId().toString(); + const chatHistory = new AzureCosmosDBMongoChatMessageHistory( + dbcfg, + sessionId + ); + + const blankResult = await chatHistory.getMessages(); + expect(blankResult).toStrictEqual([]); + + await chatHistory.addUserMessage("Who is the best vocalist?"); + await chatHistory.addAIChatMessage("Ozzy Osbourne"); + + const expectedMessages = [ + new HumanMessage("Who is the best vocalist?"), + new AIMessage("Ozzy Osbourne"), + ]; + + const resultWithHistory = await chatHistory.getMessages(); + console.log(resultWithHistory); + expect(resultWithHistory).toEqual(expectedMessages); + + await mongoClient.close(); +}); + +test("Test clear Azure Cosmos MongoDB history store", async () => { + expect(process.env.AZURE_COSMOSDB_MONGODB_CONNECTION_STRING).toBeDefined(); + + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + const mongoClient = new MongoClient( + process.env.AZURE_COSMOSDB_MONGODB_CONNECTION_STRING! + ); + const dbcfg: AzureCosmosDBMongoChatHistoryDBConfig = { + client: mongoClient, + connectionString: process.env.AZURE_COSMOSDB_MONGODB_CONNECTION_STRING, + databaseName: "langchain", + collectionName: "chathistory", + }; + + const sessionId = new ObjectId().toString(); + const chatHistory = new AzureCosmosDBMongoChatMessageHistory( + dbcfg, + sessionId + ); + + await chatHistory.addUserMessage("Who is the best vocalist?"); + await chatHistory.addAIChatMessage("Ozzy Osbourne"); + + const expectedMessages = [ + new HumanMessage("Who is the best vocalist?"), + new AIMessage("Ozzy Osbourne"), + ]; + + const resultWithHistory = await chatHistory.getMessages(); + expect(resultWithHistory).toEqual(expectedMessages); + + await chatHistory.clear(); + + const blankResult = await chatHistory.getMessages(); + expect(blankResult).toStrictEqual([]); + + await mongoClient.close(); +}); diff --git a/libs/langchain-azure-cosmosdb/src/tests/chat_histories.int.test.ts b/libs/langchain-azure-cosmosdb/src/tests/chat_histories/nosql.int.test.ts similarity index 98% rename from libs/langchain-azure-cosmosdb/src/tests/chat_histories.int.test.ts rename to libs/langchain-azure-cosmosdb/src/tests/chat_histories/nosql.int.test.ts index 76da66d7f805..9a6b12b125b6 100644 --- a/libs/langchain-azure-cosmosdb/src/tests/chat_histories.int.test.ts +++ b/libs/langchain-azure-cosmosdb/src/tests/chat_histories/nosql.int.test.ts @@ -6,7 +6,7 @@ import { HumanMessage, AIMessage } from "@langchain/core/messages"; import { CosmosClient } from "@azure/cosmos"; import { DefaultAzureCredential } from "@azure/identity"; import { ObjectId } from "mongodb"; -import { AzureCosmsosDBNoSQLChatMessageHistory } from "../chat_histories.js"; +import { AzureCosmsosDBNoSQLChatMessageHistory } from "../../chat_histories/nosql.js"; const DATABASE_NAME = "langchainTestDB"; const CONTAINER_NAME = "testContainer"; diff --git a/yarn.lock b/yarn.lock index 1a6472c33e59..9199c735c85e 100644 --- a/yarn.lock +++ b/yarn.lock @@ -13491,12 +13491,12 @@ __metadata: languageName: node linkType: hard -"@mongodb-js/saslprep@npm:^1.1.5": - version: 1.1.8 - resolution: "@mongodb-js/saslprep@npm:1.1.8" +"@mongodb-js/saslprep@npm:^1.1.9": + version: 1.1.9 + resolution: "@mongodb-js/saslprep@npm:1.1.9" dependencies: sparse-bitfield: ^3.0.3 - checksum: 259fda7ec913b5e63f102ae18840ef0d811c4a50919bbc437bd0452980806d640cd06c36076ed1655f1581ef24cd7316be0671f4b7429e7c97c7066524d2dbee + checksum: 6f13983e41c9fbd5273eeae9135e47e5b7a19125a63287bea69e33a618f8e034cfcf2258c77d0f5d6dcf386dfe2bb520bc01613afd1528c52f82c71172629242 languageName: node linkType: hard @@ -22627,6 +22627,13 @@ __metadata: languageName: node linkType: hard +"bson@npm:^6.10.0": + version: 6.10.1 + resolution: "bson@npm:6.10.1" + checksum: 7c85c8df309bbfd4d42fae54aa37112ee048a89457be908a0e53a01d077d548c94a5a6870dd725ef48130da935286edc8b9ce04830869446db22b8c13a370c42 + languageName: node + linkType: hard + "bson@npm:^6.2.0": version: 6.2.0 resolution: "bson@npm:6.2.0" @@ -22634,13 +22641,6 @@ __metadata: languageName: node linkType: hard -"bson@npm:^6.7.0": - version: 6.8.0 - resolution: "bson@npm:6.8.0" - checksum: 66076b04d7d54e7773d601a19b7c224bc5cff6b008efe102463fbc058879f2c84c0ed793b5b6ed12cc7616bbbe5e670db81cf7352e0ea947918119f8af704ba5 - languageName: node - linkType: hard - "buffer-alloc-unsafe@npm:^1.1.0": version: 1.1.0 resolution: "buffer-alloc-unsafe@npm:1.1.0" @@ -35049,11 +35049,11 @@ __metadata: linkType: hard "mongodb@npm:^6.10.0": - version: 6.10.0 - resolution: "mongodb@npm:6.10.0" + version: 6.11.0 + resolution: "mongodb@npm:6.11.0" dependencies: - "@mongodb-js/saslprep": ^1.1.5 - bson: ^6.7.0 + "@mongodb-js/saslprep": ^1.1.9 + bson: ^6.10.0 mongodb-connection-string-url: ^3.0.0 peerDependencies: "@aws-sdk/credential-providers": ^3.188.0 @@ -35078,7 +35078,7 @@ __metadata: optional: true socks: optional: true - checksum: b8e7ab9fb84181cb020b5fef5fedd90a5fc12140e688fa12ba588d523a958bb9f8790bfaceeca9f594171794eda0f56be855d7d0588705db82b3de7bf5e2352c + checksum: cb677bdee565eb9e7cbc27e538d5fafe61312e8ccfc97d4e04fdbf282f03e566537d659394f3236b6f958392707924a7ecae86fcb038a2f8f47b8edafa6edf4d languageName: node linkType: hard