feat(cosmosdbnosql): Add Chat History Integration (#7057)

docs/core_docs/docs/integrations/memory/azure_cosmosdb_nosql.mdx

@@ -0,0 +1,42 @@
+---
+hide_table_of_contents: true
+---
+
+import CodeBlock from "@theme/CodeBlock";
+
+# Azure Cosmos DB NoSQL Chat Message History
+
+The AzureCosmosDBNoSQLChatMessageHistory uses Cosmos DB to store chat message history. For longer-term persistence across chat sessions, you can swap out the default in-memory `chatHistory` that backs chat memory classes like `BufferMemory`.
+If you don't have an Azure account, you can [create a free account](https://azure.microsoft.com/free/) to get started.
+
+## Setup
+
+You'll first need to install the [`@langchain/azure-cosmosdb`](https://www.npmjs.com/package/@langchain/azure-cosmosdb) package:
+
+```bash npm2yarn
+npm install @langchain/azure-cosmosdb @langchain/core
+```
+
+import IntegrationInstallTooltip from "@mdx_components/integration_install_tooltip.mdx";
+
+<IntegrationInstallTooltip></IntegrationInstallTooltip>
+
+```bash npm2yarn
+npm install @langchain/openai @langchain/community @langchain/core
+```
+
+You'll also need to have an Azure Cosmos DB for NoSQL instance running. You can deploy a free version on Azure Portal without any cost, following [this guide](https://learn.microsoft.com/azure/cosmos-db/nosql/quickstart-portal).
+
+Once you have your instance running, make sure you have the connection string. If you are using Managed Identity, you need to have the endpoint. You can find them in the Azure Portal, under the "Settings / Keys" section of your instance.
+
+:::info
+
+When using Azure Managed Identity and role-based access control, you must ensure that the database and container have been created beforehand. RBAC does not provide permissions to create databases and containers. You can get more information about the permission model in the [Azure Cosmos DB documentation](https://learn.microsoft.com/azure/cosmos-db/how-to-setup-rbac#permission-model).
+
+:::
+
+## Usage
+
+import Example from "@examples/memory/azure_cosmosdb_nosql.ts";
+
+<CodeBlock language="typescript">{Example}</CodeBlock>

docs/core_docs/docs/integrations/platforms/microsoft.mdx

@@ -150,6 +150,22 @@ See a [usage example](/docs/integrations/llm_caching/azure_cosmosdb_nosql).
 import { AzureCosmosDBNoSQLSemanticCache } from "@langchain/azure-cosmosdb";
 ```
 
+## Chat Message History
+
+### Azure Cosmos DB NoSQL Chat Message History
+
+> The AzureCosmosDBNoSQLChatMessageHistory uses Cosmos DB to store chat message history. For longer-term persistence across chat sessions, you can swap out the default in-memory `chatHistory` that backs chat memory classes like `BufferMemory`.
+
+```bash npm2yarn
+npm install @langchain/azure-cosmosdb @langchain/core
+```
+
+See [usage example](/docs/integrations/memory/azure_cosmosdb_nosql.mdx).
+
+```typescript
+import { AzureCosmosDBNoSQLChatMessageHistory } from "@langchain/azure-cosmosdb";
+```
+
 ## Document loaders
 
 ### Azure Blob Storage

examples/src/memory/azure_cosmosdb_nosql.ts

@@ -0,0 +1,58 @@
+import { ChatOpenAI } from "@langchain/openai";
+import { AzureCosmsosDBNoSQLChatMessageHistory } from "@langchain/azure-cosmosdb";
+import { RunnableWithMessageHistory } from "@langchain/core/runnables";
+import { StringOutputParser } from "@langchain/core/output_parsers";
+import {
+  ChatPromptTemplate,
+  MessagesPlaceholder,
+} from "@langchain/core/prompts";
+
+const model = new ChatOpenAI({
+  model: "gpt-3.5-turbo",
+  temperature: 0,
+});
+
+const prompt = ChatPromptTemplate.fromMessages([
+  [
+    "system",
+    "You are a helpful assistant. Answer all questions to the best of your ability.",
+  ],
+  new MessagesPlaceholder("chat_history"),
+  ["human", "{input}"],
+]);
+
+const chain = prompt.pipe(model).pipe(new StringOutputParser());
+
+const chainWithHistory = new RunnableWithMessageHistory({
+  runnable: chain,
+  inputMessagesKey: "input",
+  historyMessagesKey: "chat_history",
+  getMessageHistory: async (sessionId) => {
+    const chatHistory = new AzureCosmsosDBNoSQLChatMessageHistory({
+      sessionId,
+      userId: "user-id",
+      databaseName: "DATABASE_NAME",
+      containerName: "CONTAINER_NAME",
+    });
+    return chatHistory;
+  },
+});
+
+const res1 = await chainWithHistory.invoke(
+  { input: "Hi! I'm Jim." },
+  { configurable: { sessionId: "langchain-test-session" } }
+);
+console.log({ res1 });
+/*
+{ res1: 'Hi Jim! How can I assist you today?' }
+ */
+
+const res2 = await chainWithHistory.invoke(
+  { input: "What did I just say my name was?" },
+  { configurable: { sessionId: "langchain-test-session" } }
+);
+console.log({ res2 });
+
+/*
+ { res2: { response: 'You said your name was Jim.' } 
+  */

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

@@ -0,0 +1,204 @@
+import { Container, CosmosClient, CosmosClientOptions } from "@azure/cosmos";
+import { DefaultAzureCredential, TokenCredential } from "@azure/identity";
+import { BaseListChatMessageHistory } from "@langchain/core/chat_history";
+import {
+  BaseMessage,
+  mapChatMessagesToStoredMessages,
+  mapStoredMessagesToChatMessages,
+} from "@langchain/core/messages";
+import { getEnvironmentVariable } from "@langchain/core/utils/env";
+
+const USER_AGENT_SUFFIX = "langchainjs-cdbnosql-chathistory-javascript";
+const DEFAULT_DATABASE_NAME = "chatHistoryDB";
+const DEFAULT_CONTAINER_NAME = "chatHistoryContainer";
+
+/**
+ * Type for the input to the `AzureCosmosDBNoSQLChatMessageHistory` constructor.
+ */
+export interface AzureCosmosDBNoSQLChatMessageHistoryInput {
+  sessionId: string;
+  userId?: string;
+  client?: CosmosClient;
+  connectionString?: string;
+  endpoint?: string;
+  databaseName?: string;
+  containerName?: string;
+  credentials?: TokenCredential;
+  ttl?: number;
+}
+
+/**
+ * Class for storing chat message history with Cosmos DB NoSQL. It extends the
+ * BaseListChatMessageHistory class and provides methods to get, add, and
+ * clear messages.
+ *
+ * @example
+ * ```typescript
+ *  const model = new ChatOpenAI({
+ *   model: "gpt-3.5-turbo",
+ *   temperature: 0,
+ * });
+ * const prompt = ChatPromptTemplate.fromMessages([
+ *   [
+ *     "system",
+ *     "You are a helpful assistant. Answer all questions to the best of your ability.",
+ *   ],
+ *   new MessagesPlaceholder("chat_history"),
+ *   ["human", "{input}"],
+ * ]);
+ *
+ * const chain = prompt.pipe(model).pipe(new StringOutputParser());
+ * const chainWithHistory = new RunnableWithMessageHistory({
+ *   runnable: chain,
+ *  inputMessagesKey: "input",
+ *  historyMessagesKey: "chat_history",
+ *   getMessageHistory: async (sessionId) => {
+ *     const chatHistory = new AzureCosmsosDBNoSQLChatMessageHistory({
+ *       sessionId: sessionId,
+ *       userId: "user-id",
+ *       databaseName: "DATABASE_NAME",
+ *       containerName: "CONTAINER_NAME",
+ *     })
+ *     return chatHistory;
+ *   },
+ * });
+ * await chainWithHistory.invoke(
+ *   { input: "What did I just say my name was?" },
+ *   { configurable: { sessionId: "session-id" } }
+ * );
+ * ```
+ */
+
+export class AzureCosmsosDBNoSQLChatMessageHistory extends BaseListChatMessageHistory {
+  lc_namespace = ["langchain", "stores", "message", "azurecosmosdb"];
+
+  private container: Container;
+
+  private sessionId: string;
+
+  private databaseName: string;
+
+  private containerName: string;
+
+  private client: CosmosClient;
+
+  private userId: string;
+
+  private ttl: number | undefined;
+
+  private messageList: BaseMessage[] = [];
+
+  private initPromise?: Promise<void>;
+
+  constructor(chatHistoryInput: AzureCosmosDBNoSQLChatMessageHistoryInput) {
+    super();
+
+    this.sessionId = chatHistoryInput.sessionId;
+    this.databaseName = chatHistoryInput.databaseName ?? DEFAULT_DATABASE_NAME;
+    this.containerName =
+      chatHistoryInput.containerName ?? DEFAULT_CONTAINER_NAME;
+    this.userId = chatHistoryInput.userId ?? "anonymous";
+    this.ttl = chatHistoryInput.ttl;
+    this.client = this.initializeClient(chatHistoryInput);
+  }
+
+  private initializeClient(
+    input: AzureCosmosDBNoSQLChatMessageHistoryInput
+  ): CosmosClient {
+    const connectionString =
+      input.connectionString ??
+      getEnvironmentVariable("AZURE_COSMOSDB_NOSQL_CONNECTION_STRING");
+    const endpoint =
+      input.endpoint ?? getEnvironmentVariable("AZURE_COSMOSDB_NOSQL_ENDPOINT");
+
+    if (!input.client && !connectionString && !endpoint) {
+      throw new Error(
+        "CosmosClient, connection string, or endpoint must be provided."
+      );
+    }
+
+    if (input.client) {
+      return input.client;
+    }
+
+    if (connectionString) {
+      const [endpointPart, keyPart] = connectionString.split(";");
+      const endpoint = endpointPart.split("=")[1];
+      const key = keyPart.split("=")[1];
+
+      return new CosmosClient({
+        endpoint,
+        key,
+        userAgentSuffix: USER_AGENT_SUFFIX,
+      });
+    } else {
+      return new CosmosClient({
+        endpoint,
+        aadCredentials: input.credentials ?? new DefaultAzureCredential(),
+        userAgentSuffix: USER_AGENT_SUFFIX,
+      } as CosmosClientOptions);
+    }
+  }
+
+  private async initializeContainer(): Promise<void> {
+    if (!this.initPromise) {
+      this.initPromise = (async () => {
+        const { database } = await this.client.databases.createIfNotExists({
+          id: this.databaseName,
+        });
+        const { container } = await database.containers.createIfNotExists({
+          id: this.containerName,
+          partitionKey: "/userId",
+          defaultTtl: this.ttl,
+        });
+        this.container = container;
+      })().catch((error) => {
+        console.error("Error initializing Cosmos DB container:", error);
+        throw error;
+      });
+    }
+    return this.initPromise;
+  }
+
+  async getMessages(): Promise<BaseMessage[]> {
+    await this.initializeContainer();
+    const document = await this.container
+      .item(this.sessionId, this.userId)
+      .read();
+    const messages = document.resource?.messages || [];
+    this.messageList = mapStoredMessagesToChatMessages(messages);
+    return this.messageList;
+  }
+
+  async addMessage(message: BaseMessage): Promise<void> {
+    await this.initializeContainer();
+    this.messageList = await this.getMessages();
+    this.messageList.push(message);
+    const messages = mapChatMessagesToStoredMessages(this.messageList);
+    await this.container.items.upsert({
+      id: this.sessionId,
+      userId: this.userId,
+      messages,
+    });
+  }
+
+  async clear(): Promise<void> {
+    this.messageList = [];
+    await this.initializeContainer();
+    await this.container.item(this.sessionId, this.userId).delete();
+  }
+
+  async clearAllSessionsForUser(userId: string) {
+    await this.initializeContainer();
+    const query = {
+      query: "SELECT c.id FROM c WHERE c.userId = @userId",
+      parameters: [{ name: "@userId", value: userId }],
+    };
+    const { resources: userSessions } = await this.container.items
+      .query(query)
+      .fetchAll();
+    for (const userSession of userSessions) {
+      await this.container.item(userSession.id, userId).delete();
+    }
+  }
+}

libs/langchain-azure-cosmosdb/src/index.ts

@@ -1,3 +1,4 @@
 export * from "./azure_cosmosdb_mongodb.js";
 export * from "./azure_cosmosdb_nosql.js";
 export * from "./caches.js";
+export * from "./chat_histories.js";