resolve comments

libs/langchain-azure-cosmosdb/src/chat_histories/mongodb.ts

@@ -0,0 +1,157 @@
+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;
+}
+
+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<void>;
+
+  private readonly client: MongoClient | undefined;
+
+  private database: Db;
+
+  private collection: Collection<AzureCosmosMongoDBDocument>;
+
+  private sessionId: string;
+
+  private idKey = "sessionId";
+
+  initialize: () => Promise<void>;
+
+  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 ?? "documentsDB";
+    const collectionName = dbConfig.collectionName ?? "documents";
+
+    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<void> {
+    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<BaseMessage[]> {
+    await this.initialize();
+
+    const document = await this.collection.findOne({
+      [this.idKey]: 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<void> {
+    await this.initialize();
+
+    const messages = mapChatMessagesToStoredMessages([message]);
+    await this.collection.updateOne(
+      { [this.idKey]: this.sessionId },
+      {
+        $push: { messages: { $each: messages } } as PushOperator<Document>,
+      },
+      { upsert: true }
+    );
+  }
+
+  /**
+   * Clear the history.
+   * @returns A promise that resolves when the history has been cleared.
+   */
+  async clear(): Promise<void> {
+    await this.initialize();
+
+    await this.collection.deleteOne({ [this.idKey]: this.sessionId });
+  }
+}

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";

libs/langchain-azure-cosmosdb/src/tests/chat_histories.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";

libs/langchain-azure-cosmosdb/src/tests/chat_histories_azure_cosmosdb_mongodb.int.test.ts

@@ -2,7 +2,10 @@
 
 import { MongoClient, ObjectId } from "mongodb";
 import { AIMessage, HumanMessage } from "@langchain/core/messages";
-import { AzureCosmosDBMongoChatMessageHistory } from "../chat_histories_azure_cosmosdb_mongodb.js";
+import {
+  AzureCosmosDBMongoChatMessageHistory,
+  AzureCosmosDBMongoChatHistoryDBConfig,
+} from "../chat_histories/mongodb.js";
 
 afterAll(async () => {
   // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
@@ -18,17 +21,21 @@ 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 client = new MongoClient(
+  const mongoClient = new MongoClient(
     process.env.AZURE_COSMOSDB_MONGODB_CONNECTION_STRING!
   );
-  await client.connect();
-  const collection = client.db("langchain").collection("memory");
+  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({
-    collection,
-    sessionId,
-  });
+  const chatHistory = new AzureCosmosDBMongoChatMessageHistory(
+    dbcfg,
+    sessionId
+  );
 
   const blankResult = await chatHistory.getMessages();
   expect(blankResult).toStrictEqual([]);
@@ -45,24 +52,28 @@ test("Test Azure Cosmos MongoDB history store", async () => {
   console.log(resultWithHistory);
   expect(resultWithHistory).toEqual(expectedMessages);
 
-  await client.close();
+  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 client = new MongoClient(
+  const mongoClient = new MongoClient(
     process.env.AZURE_COSMOSDB_MONGODB_CONNECTION_STRING!
   );
-  await client.connect();
-  const collection = client.db("langchain").collection("memory");
+  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({
-    collection,
-    sessionId,
-  });
+  const chatHistory = new AzureCosmosDBMongoChatMessageHistory(
+    dbcfg,
+    sessionId
+  );
 
   await chatHistory.addUserMessage("Who is the best vocalist?");
   await chatHistory.addAIChatMessage("Ozzy Osbourne");
@@ -80,5 +91,5 @@ test("Test clear Azure Cosmos MongoDB history store", async () => {
   const blankResult = await chatHistory.getMessages();
   expect(blankResult).toStrictEqual([]);
 
-  await client.close();
+  await mongoClient.close();
 });