add session context for a user mongodb

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

@@ -20,6 +20,11 @@ export interface AzureCosmosDBMongoChatHistoryDBConfig {
   readonly collectionName?: string;
 }
 
+export type ChatSession = {
+  id: string;
+  context: Record<string, unknown>;
+};
+
 const ID_KEY = "sessionId";
 
 export class AzureCosmosDBMongoChatMessageHistory extends BaseListChatMessageHistory {
@@ -33,6 +38,8 @@ export class AzureCosmosDBMongoChatMessageHistory extends BaseListChatMessageHis
 
   private initPromise?: Promise<void>;
 
+  private context: Record<string, unknown> = {};
+
   private readonly client: MongoClient | undefined;
 
   private database: Db;
@@ -134,10 +141,12 @@ export class AzureCosmosDBMongoChatMessageHistory extends BaseListChatMessageHis
     await this.initialize();
 
     const messages = mapChatMessagesToStoredMessages([message]);
+    const context = await this.getContext();
     await this.collection.updateOne(
       { [ID_KEY]: this.sessionId },
       {
         $push: { messages: { $each: messages } } as PushOperator<Document>,
+        $set: { context },
       },
       { upsert: true }
     );
@@ -152,4 +161,51 @@ export class AzureCosmosDBMongoChatMessageHistory extends BaseListChatMessageHis
 
     await this.collection.deleteOne({ [ID_KEY]: this.sessionId });
   }
+
+  async getAllSessions(): Promise<ChatSession[]> {
+    await this.initialize();
+    const documents = await this.collection.find().toArray();
+
+    const chatSessions: ChatSession[] = documents.map((doc) => ({
+      id: doc[ID_KEY],
+      context: doc.context || {},
+    }));
+
+    return chatSessions;
+  }
+
+  async clearAllSessions() {
+    await this.initialize();
+    try {
+      await this.collection.deleteMany({});
+    } catch (error) {
+      console.log("Error clearing sessions:", error);
+    }
+  }
+
+  async getContext(): Promise<Record<string, unknown>> {
+    await this.initialize();
+
+    const document = await this.collection.findOne({
+      [ID_KEY]: this.sessionId,
+    });
+    this.context = document?.context || this.context;
+    return this.context;
+  }
+
+  async setContext(context: Record<string, unknown>): Promise<void> {
+    await this.initialize();
+
+    try {
+      await this.collection.updateOne(
+        { [ID_KEY]: this.sessionId },
+        {
+          $set: { context },
+        },
+        { upsert: true }
+      );
+    } catch (error) {
+      console.log("Error setting context", error);
+    }
+  }
 }

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

@@ -93,3 +93,46 @@ test("Test clear Azure Cosmos MongoDB history store", async () => {
 
   await mongoClient.close();
 });
+
+test("Test getAllSessions and clearAllSessions", 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 sessionId1 = new ObjectId().toString();
+  const sessionId2 = new ObjectId().toString();
+
+  const chatHistory1 = new AzureCosmosDBMongoChatMessageHistory(
+    dbcfg,
+    sessionId1
+  );
+  const chatHistory2 = new AzureCosmosDBMongoChatMessageHistory(
+    dbcfg,
+    sessionId2
+  );
+
+  await chatHistory1.addUserMessage("What is AI?");
+  await chatHistory1.addAIChatMessage("AI stands for Artificial Intelligence.");
+  await chatHistory2.addUserMessage("What is the best programming language?");
+  await chatHistory2.addAIChatMessage("It depends on the use case.");
+
+  const allSessions = await chatHistory1.getAllSessions();
+  expect(allSessions.length).toBe(2);
+  expect(allSessions[0].id).toBe(sessionId1);
+  expect(allSessions[1].id).toBe(sessionId2);
+
+  await chatHistory1.clearAllSessions();
+  const clearedSessions = await chatHistory1.getAllSessions();
+  expect(clearedSessions.length).toBe(0);
+
+  await mongoClient.close();
+});