Skip to content

Commit

Permalink
mongodb[patch]: Migrate MongoDBChatMessageHistory (langchain-ai#18590)
Browse files Browse the repository at this point in the history
## **Description** 
Migrate the `MongoDBChatMessageHistory` to the managed
`langchain-mongodb` partner-package
## **Dependencies**
None
## **Twitter handle**
@mongodb

## **tests and docs**
- [x] Migrate existing integration test
- [x ]~ Convert existing integration test to a unit test~ Creation is
out of scope for this ticket
- [x ] ~Considering delaying work until langchain-ai#17470 merges to leverage the
`MockCollection` object. ~
- [x] **Lint and test**: Run `make format`, `make lint` and `make test`
from the root of the package(s) you've modified. See contribution
guidelines for more: https://python.langchain.com/docs/contributing/

---------

Co-authored-by: Erick Friis <erick@langchain.dev>
  • Loading branch information
2 people authored and Dave Bechberger committed Mar 29, 2024
1 parent c87f267 commit 324c59d
Show file tree
Hide file tree
Showing 8 changed files with 858 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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",
"```"
]
},
Expand Down Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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.
Expand Down
2 changes: 2 additions & 0 deletions libs/partners/mongodb/langchain_mongodb/__init__.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
from langchain_mongodb.chat_message_histories import MongoDBChatMessageHistory
from langchain_mongodb.vectorstores import (
MongoDBAtlasVectorSearch,
)

__all__ = [
"MongoDBAtlasVectorSearch",
"MongoDBChatMessageHistory",
]
84 changes: 84 additions & 0 deletions libs/partners/mongodb/langchain_mongodb/chat_message_histories.py
Original file line number Diff line number Diff line change
@@ -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)
726 changes: 724 additions & 2 deletions libs/partners/mongodb/poetry.lock

Large diffs are not rendered by default.

1 change: 1 addition & 0 deletions libs/partners/mongodb/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
Original file line number Diff line number Diff line change
@@ -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 == []
4 changes: 1 addition & 3 deletions libs/partners/mongodb/tests/unit_tests/test_imports.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
from langchain_mongodb import __all__

EXPECTED_ALL = [
"MongoDBAtlasVectorSearch",
]
EXPECTED_ALL = ["MongoDBAtlasVectorSearch", "MongoDBChatMessageHistory"]


def test_all_imports() -> None:
Expand Down

0 comments on commit 324c59d

Please sign in to comment.