langchain-mongodb: Migrate MongoDBChatMessageHistory

docs/docs/integrations/memory/mongodb_chat_message_history.ipynb

@@ -21,10 +21,10 @@
    "source": [
     "## Setup\n",
     "\n",
-    "The integration lives in the `langchain-community` package, so we need to install that. We also need to install the `pymongo` package.\n",
+    "The integration lives in the `langchain-mongodb` package, so we need to install that.\n",
     "\n",
     "```bash\n",
-    "pip install -U --quiet langchain-community pymongo\n",
+    "pip install -U --quiet langchain-mongodb\n",
     "```"
    ]
   },
@@ -80,7 +80,7 @@
    },
    "outputs": [],
    "source": [
-    "from langchain_community.chat_message_histories import MongoDBChatMessageHistory\n",
+    "from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory\n",
     "\n",
     "chat_message_history = MongoDBChatMessageHistory(\n",
     "    session_id=\"test_session\",\n",

libs/community/langchain_community/chat_message_histories/mongodb.py

@@ -2,6 +2,7 @@
 import logging
 from typing import List
 
+from langchain_core._api.deprecation import deprecated
 from langchain_core.chat_history import BaseChatMessageHistory
 from langchain_core.messages import (
     BaseMessage,
@@ -15,6 +16,11 @@
 DEFAULT_COLLECTION_NAME = "message_store"
 
 
+@deprecated(
+    since="0.0.25",
+    removal="0.2.0",
+    alternative_import="langchain_mongodb.MongoDBChatMessageHistory",
+)
 class MongoDBChatMessageHistory(BaseChatMessageHistory):
     """Chat message history that stores history in MongoDB.
 

libs/partners/mongodb/langchain_mongodb/__init__.py

@@ -1,7 +1,9 @@
+from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory
 from langchain_mongodb.vectorstores import (
     MongoDBAtlasVectorSearch,
 )
 
 __all__ = [
     "MongoDBAtlasVectorSearch",
+    "MongoDBChatMessageHistory",
 ]

libs/partners/mongodb/langchain_mongodb/chat_message_histories.py

@@ -0,0 +1,84 @@
+import json
+import logging
+from typing import List
+
+from langchain_core.chat_history import BaseChatMessageHistory
+from langchain_core.messages import (
+    BaseMessage,
+    message_to_dict,
+    messages_from_dict,
+)
+from pymongo import MongoClient, errors
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_DBNAME = "chat_history"
+DEFAULT_COLLECTION_NAME = "message_store"
+
+
+class MongoDBChatMessageHistory(BaseChatMessageHistory):
+    """Chat message history that stores history in MongoDB.
+
+    Args:
+        connection_string: connection string to connect to MongoDB
+        session_id: arbitrary key that is used to store the messages
+            of a single chat session.
+        database_name: name of the database to use
+        collection_name: name of the collection to use
+    """
+
+    def __init__(
+        self,
+        connection_string: str,
+        session_id: str,
+        database_name: str = DEFAULT_DBNAME,
+        collection_name: str = DEFAULT_COLLECTION_NAME,
+    ):
+        self.connection_string = connection_string
+        self.session_id = session_id
+        self.database_name = database_name
+        self.collection_name = collection_name
+
+        try:
+            self.client: MongoClient = MongoClient(connection_string)
+        except errors.ConnectionFailure as error:
+            logger.error(error)
+
+        self.db = self.client[database_name]
+        self.collection = self.db[collection_name]
+        self.collection.create_index("SessionId")
+
+    @property
+    def messages(self) -> List[BaseMessage]:  # type: ignore
+        """Retrieve the messages from MongoDB"""
+        try:
+            cursor = self.collection.find({"SessionId": self.session_id})
+        except errors.OperationFailure as error:
+            logger.error(error)
+
+        if cursor:
+            items = [json.loads(document["History"]) for document in cursor]
+        else:
+            items = []
+
+        messages = messages_from_dict(items)
+        return messages
+
+    def add_message(self, message: BaseMessage) -> None:
+        """Append the message to the record in MongoDB"""
+        try:
+            self.collection.insert_one(
+                {
+                    "SessionId": self.session_id,
+                    "History": json.dumps(message_to_dict(message)),
+                }
+            )
+        except errors.WriteError as err:
+            logger.error(err)
+
+    def clear(self) -> None:
+        """Clear session memory from MongoDB"""
+        try:
+            self.collection.delete_many({"SessionId": self.session_id})
+        except errors.WriteError as err:
+            logger.error(err)

libs/partners/mongodb/pyproject.toml

@@ -26,6 +26,7 @@ pytest-mock = "^3.10.0"
 syrupy = "^4.0.2"
 pytest-watcher = "^0.3.4"
 pytest-asyncio = "^0.21.1"
+langchain = { path = "../../langchain", develop = true }
 langchain-core = { path = "../../core", develop = true }
 
 [tool.poetry.group.codespell]

libs/partners/mongodb/tests/integration_tests/test_chat_message_histories.py

@@ -0,0 +1,37 @@
+import json
+import os
+
+from langchain.memory import ConversationBufferMemory
+from langchain_core.messages import message_to_dict
+
+from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory
+
+# Replace these with your mongodb connection string
+connection_string = os.environ.get("MONGODB_ATLAS_URI", "")
+
+
+def test_memory_with_message_store() -> None:
+    """Test the memory with a message store."""
+    # setup MongoDB as a message store
+    message_history = MongoDBChatMessageHistory(
+        connection_string=connection_string, session_id="test-session"
+    )
+    memory = ConversationBufferMemory(
+        memory_key="baz", chat_memory=message_history, return_messages=True
+    )
+
+    # add some messages
+    memory.chat_memory.add_ai_message("This is me, the AI")
+    memory.chat_memory.add_user_message("This is me, the human")
+
+    # get the message history from the memory store and turn it into a json
+    messages = memory.chat_memory.messages
+    messages_json = json.dumps([message_to_dict(msg) for msg in messages])
+
+    assert "This is me, the AI" in messages_json
+    assert "This is me, the human" in messages_json
+
+    # remove the record from Azure Cosmos DB, so the next test run won't pick it up
+    memory.chat_memory.clear()
+
+    assert memory.chat_memory.messages == []

libs/partners/mongodb/tests/unit_tests/test_imports.py

@@ -1,8 +1,6 @@
 from langchain_mongodb import __all__
 
-EXPECTED_ALL = [
-    "MongoDBAtlasVectorSearch",
-]
+EXPECTED_ALL = ["MongoDBAtlasVectorSearch", "MongoDBChatMessageHistory"]
 
 
 def test_all_imports() -> None: